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

use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::Loadable;
use blobstore::LoadableError;
use blobstore::Storable;
use context::CoreContext;
use futures::stream;
use futures::stream::BoxStream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use mononoke_types::basename_suffix_skeleton_manifest::BasenameSuffixSkeletonManifest;
use mononoke_types::basename_suffix_skeleton_manifest::BssmDirectory;
use mononoke_types::basename_suffix_skeleton_manifest::BssmEntry;
use mononoke_types::fsnode::Fsnode;
use mononoke_types::fsnode::FsnodeEntry;
use mononoke_types::fsnode::FsnodeFile;
use mononoke_types::skeleton_manifest::SkeletonManifest;
use mononoke_types::skeleton_manifest::SkeletonManifestEntry;
use mononoke_types::unode::ManifestUnode;
use mononoke_types::unode::UnodeEntry;
use mononoke_types::FileUnodeId;
use mononoke_types::FsnodeId;
use mononoke_types::MPath;
use mononoke_types::MPathElement;
use mononoke_types::ManifestUnodeId;
use mononoke_types::SkeletonManifestId;
use serde_derive::Deserialize;
use serde_derive::Serialize;

#[async_trait]
pub trait AsyncManifest<Store: Send + Sync>: Sized + 'static {
    type TreeId: Send + Sync;
    type LeafId: Send + Sync;

    async fn list(
        &self,
        ctx: &CoreContext,
        blobstore: &Store,
    ) -> Result<BoxStream<'async_trait, Result<(MPathElement, Entry<Self::TreeId, Self::LeafId>)>>>;
    /// List all subentries with a given prefix
    async fn list_prefix(
        &self,
        ctx: &CoreContext,
        blobstore: &Store,
        prefix: &[u8],
    ) -> Result<BoxStream<'async_trait, Result<(MPathElement, Entry<Self::TreeId, Self::LeafId>)>>>;
    async fn lookup(
        &self,
        ctx: &CoreContext,
        blobstore: &Store,
        name: &MPathElement,
    ) -> Result<Option<Entry<Self::TreeId, Self::LeafId>>>;
}

pub trait Manifest: Sync + Sized + 'static {
    type TreeId: Send + Sync;
    type LeafId: Send + Sync;
    fn list(&self) -> Box<dyn Iterator<Item = (MPathElement, Entry<Self::TreeId, Self::LeafId>)>>;
    /// List all subentries with a given prefix
    fn list_prefix<'a>(
        &'a self,
        prefix: &'a [u8],
    ) -> Box<dyn Iterator<Item = (MPathElement, Entry<Self::TreeId, Self::LeafId>)> + 'a> {
        Box::new(self.list().filter(|(k, _)| k.starts_with(prefix)))
    }
    fn lookup(&self, name: &MPathElement) -> Option<Entry<Self::TreeId, Self::LeafId>>;
}

#[async_trait]
impl<M: Manifest, Store: Send + Sync> AsyncManifest<Store> for M {
    type TreeId = <Self as Manifest>::TreeId;
    type LeafId = <Self as Manifest>::LeafId;

    async fn list(
        &self,
        _ctx: &CoreContext,
        _blobstore: &Store,
    ) -> Result<BoxStream<'async_trait, Result<(MPathElement, Entry<Self::TreeId, Self::LeafId>)>>>
    {
        Ok(stream::iter(Manifest::list(self).map(anyhow::Ok).collect::<Vec<_>>()).boxed())
    }

    async fn list_prefix(
        &self,
        _ctx: &CoreContext,
        _blobstore: &Store,
        prefix: &[u8],
    ) -> Result<BoxStream<'async_trait, Result<(MPathElement, Entry<Self::TreeId, Self::LeafId>)>>>
    {
        Ok(stream::iter(
            Manifest::list_prefix(self, prefix)
                .map(anyhow::Ok)
                .collect::<Vec<_>>(),
        )
        .boxed())
    }

    async fn lookup(
        &self,
        _ctx: &CoreContext,
        _blobstore: &Store,
        name: &MPathElement,
    ) -> Result<Option<Entry<Self::TreeId, Self::LeafId>>> {
        anyhow::Ok(Manifest::lookup(self, name))
    }
}

fn to_mf_entry(entry: BssmEntry) -> Entry<BssmDirectory, ()> {
    match entry {
        BssmEntry::Directory(dir) => Entry::Tree(dir),
        BssmEntry::File => Entry::Leaf(()),
    }
}

#[async_trait]
impl<Store: Blobstore> AsyncManifest<Store> for BasenameSuffixSkeletonManifest {
    type TreeId = BssmDirectory;
    type LeafId = ();

    async fn list(
        &self,
        ctx: &CoreContext,
        blobstore: &Store,
    ) -> Result<BoxStream<'async_trait, Result<(MPathElement, Entry<Self::TreeId, Self::LeafId>)>>>
    {
        anyhow::Ok(
            self.clone()
                .into_subentries(ctx, blobstore)
                .map_ok(|(path, entry)| (path, to_mf_entry(entry)))
                .boxed(),
        )
    }

    async fn list_prefix(
        &self,
        ctx: &CoreContext,
        blobstore: &Store,
        prefix: &[u8],
    ) -> Result<BoxStream<'async_trait, Result<(MPathElement, Entry<Self::TreeId, Self::LeafId>)>>>
    {
        anyhow::Ok(
            self.clone()
                .into_prefix_subentries(ctx, blobstore, prefix)
                .map_ok(|(path, entry)| (path, to_mf_entry(entry)))
                .boxed(),
        )
    }

    async fn lookup(
        &self,
        ctx: &CoreContext,
        blobstore: &Store,
        name: &MPathElement,
    ) -> Result<Option<Entry<Self::TreeId, Self::LeafId>>> {
        Ok(self.lookup(ctx, blobstore, name).await?.map(to_mf_entry))
    }
}

impl Manifest for ManifestUnode {
    type TreeId = ManifestUnodeId;
    type LeafId = FileUnodeId;

    fn lookup(&self, name: &MPathElement) -> Option<Entry<Self::TreeId, Self::LeafId>> {
        self.lookup(name).map(convert_unode)
    }

    fn list(&self) -> Box<dyn Iterator<Item = (MPathElement, Entry<Self::TreeId, Self::LeafId>)>> {
        let v: Vec<_> = self
            .list()
            .map(|(basename, entry)| (basename.clone(), convert_unode(entry)))
            .collect();
        Box::new(v.into_iter())
    }
}

fn convert_unode(unode_entry: &UnodeEntry) -> Entry<ManifestUnodeId, FileUnodeId> {
    match unode_entry {
        UnodeEntry::File(file_unode_id) => Entry::Leaf(file_unode_id.clone()),
        UnodeEntry::Directory(mf_unode_id) => Entry::Tree(mf_unode_id.clone()),
    }
}

impl Manifest for Fsnode {
    type TreeId = FsnodeId;
    type LeafId = FsnodeFile;

    fn lookup(&self, name: &MPathElement) -> Option<Entry<Self::TreeId, Self::LeafId>> {
        self.lookup(name).map(convert_fsnode)
    }

    fn list(&self) -> Box<dyn Iterator<Item = (MPathElement, Entry<Self::TreeId, Self::LeafId>)>> {
        let v: Vec<_> = self
            .list()
            .map(|(basename, entry)| (basename.clone(), convert_fsnode(entry)))
            .collect();
        Box::new(v.into_iter())
    }
}

fn convert_fsnode(fsnode_entry: &FsnodeEntry) -> Entry<FsnodeId, FsnodeFile> {
    match fsnode_entry {
        FsnodeEntry::File(fsnode_file) => Entry::Leaf(*fsnode_file),
        FsnodeEntry::Directory(fsnode_directory) => Entry::Tree(fsnode_directory.id().clone()),
    }
}

impl Manifest for SkeletonManifest {
    type TreeId = SkeletonManifestId;
    type LeafId = ();

    fn lookup(&self, name: &MPathElement) -> Option<Entry<Self::TreeId, Self::LeafId>> {
        self.lookup(name).map(convert_skeleton_manifest)
    }

    fn list(&self) -> Box<dyn Iterator<Item = (MPathElement, Entry<Self::TreeId, Self::LeafId>)>> {
        let v: Vec<_> = self
            .list()
            .map(|(basename, entry)| (basename.clone(), convert_skeleton_manifest(entry)))
            .collect();
        Box::new(v.into_iter())
    }
}

fn convert_skeleton_manifest(
    skeleton_entry: &SkeletonManifestEntry,
) -> Entry<SkeletonManifestId, ()> {
    match skeleton_entry {
        SkeletonManifestEntry::File => Entry::Leaf(()),
        SkeletonManifestEntry::Directory(skeleton_directory) => {
            Entry::Tree(skeleton_directory.id().clone())
        }
    }
}

pub type Weight = usize;

pub trait OrderedManifest: Manifest {
    fn lookup_weighted(
        &self,
        name: &MPathElement,
    ) -> Option<Entry<(Weight, <Self as Manifest>::TreeId), <Self as Manifest>::LeafId>>;
    fn list_weighted(
        &self,
    ) -> Box<
        dyn Iterator<
            Item = (
                MPathElement,
                Entry<(Weight, <Self as Manifest>::TreeId), <Self as Manifest>::LeafId>,
            ),
        >,
    >;
}

#[async_trait]
pub trait AsyncOrderedManifest<Store: Send + Sync>: AsyncManifest<Store> {
    async fn list_weighted(
        &self,
        ctx: &CoreContext,
        blobstore: &Store,
    ) -> Result<
        BoxStream<
            'async_trait,
            Result<(MPathElement, Entry<(Weight, Self::TreeId), Self::LeafId>)>,
        >,
    >;
    async fn lookup_weighted(
        &self,
        ctx: &CoreContext,
        blobstore: &Store,
        name: &MPathElement,
    ) -> Result<Option<Entry<(Weight, Self::TreeId), Self::LeafId>>>;
}

#[async_trait]
impl<M: OrderedManifest, Store: Send + Sync> AsyncOrderedManifest<Store> for M {
    async fn list_weighted(
        &self,
        _ctx: &CoreContext,
        _blobstore: &Store,
    ) -> Result<
        BoxStream<
            'async_trait,
            Result<(MPathElement, Entry<(Weight, Self::TreeId), Self::LeafId>)>,
        >,
    > {
        Ok(stream::iter(
            OrderedManifest::list_weighted(self)
                .map(anyhow::Ok)
                .collect::<Vec<_>>(),
        )
        .boxed())
    }
    async fn lookup_weighted(
        &self,
        _ctx: &CoreContext,
        _blobstore: &Store,
        name: &MPathElement,
    ) -> Result<Option<Entry<(Weight, Self::TreeId), Self::LeafId>>> {
        anyhow::Ok(OrderedManifest::lookup_weighted(self, name))
    }
}

fn convert_bssm_to_weighted(entry: Entry<BssmDirectory, ()>) -> Entry<(Weight, BssmDirectory), ()> {
    match entry {
        Entry::Tree(dir) => Entry::Tree((dir.rollup_count.try_into().unwrap_or(usize::MAX), dir)),
        Entry::Leaf(()) => Entry::Leaf(()),
    }
}

#[async_trait]
impl<Store: Blobstore> AsyncOrderedManifest<Store> for BasenameSuffixSkeletonManifest {
    async fn list_weighted(
        &self,
        ctx: &CoreContext,
        blobstore: &Store,
    ) -> Result<
        BoxStream<
            'async_trait,
            Result<(MPathElement, Entry<(Weight, Self::TreeId), Self::LeafId>)>,
        >,
    > {
        self.list(ctx, blobstore).await.map(|stream| {
            stream
                .map_ok(|(p, entry)| (p, convert_bssm_to_weighted(entry)))
                .boxed()
        })
    }

    async fn lookup_weighted(
        &self,
        ctx: &CoreContext,
        blobstore: &Store,
        name: &MPathElement,
    ) -> Result<Option<Entry<(Weight, Self::TreeId), Self::LeafId>>> {
        AsyncManifest::lookup(self, ctx, blobstore, name)
            .await
            .map(|opt| opt.map(convert_bssm_to_weighted))
    }
}

impl OrderedManifest for SkeletonManifest {
    fn lookup_weighted(
        &self,
        name: &MPathElement,
    ) -> Option<Entry<(Weight, <Self as Manifest>::TreeId), <Self as Manifest>::LeafId>> {
        self.lookup(name).map(convert_skeleton_manifest_weighted)
    }

    fn list_weighted(
        &self,
    ) -> Box<
        dyn Iterator<
            Item = (
                MPathElement,
                Entry<(Weight, <Self as Manifest>::TreeId), <Self as Manifest>::LeafId>,
            ),
        >,
    > {
        let v: Vec<_> = self
            .list()
            .map(|(basename, entry)| (basename.clone(), convert_skeleton_manifest_weighted(entry)))
            .collect();
        Box::new(v.into_iter())
    }
}

fn convert_skeleton_manifest_weighted(
    skeleton_entry: &SkeletonManifestEntry,
) -> Entry<(Weight, SkeletonManifestId), ()> {
    match skeleton_entry {
        SkeletonManifestEntry::File => Entry::Leaf(()),
        SkeletonManifestEntry::Directory(skeleton_directory) => {
            let summary = skeleton_directory.summary();
            let weight = summary.descendant_files_count + summary.descendant_dirs_count;
            Entry::Tree((weight as Weight, skeleton_directory.id().clone()))
        }
    }
}

impl OrderedManifest for Fsnode {
    fn lookup_weighted(
        &self,
        name: &MPathElement,
    ) -> Option<Entry<(Weight, <Self as Manifest>::TreeId), <Self as Manifest>::LeafId>> {
        self.lookup(name).map(convert_fsnode_weighted)
    }

    fn list_weighted(
        &self,
    ) -> Box<
        dyn Iterator<
            Item = (
                MPathElement,
                Entry<(Weight, <Self as Manifest>::TreeId), <Self as Manifest>::LeafId>,
            ),
        >,
    > {
        let v: Vec<_> = self
            .list()
            .map(|(basename, entry)| (basename.clone(), convert_fsnode_weighted(entry)))
            .collect();
        Box::new(v.into_iter())
    }
}

fn convert_fsnode_weighted(fsnode_entry: &FsnodeEntry) -> Entry<(Weight, FsnodeId), FsnodeFile> {
    match fsnode_entry {
        FsnodeEntry::File(fsnode_file) => Entry::Leaf(*fsnode_file),
        FsnodeEntry::Directory(fsnode_directory) => {
            let summary = fsnode_directory.summary();
            // Fsnodes don't have a full descendant dirs count, so we use the
            // child count as a lower-bound estimate.
            let weight = summary.descendant_files_count + summary.child_dirs_count;
            Entry::Tree((weight as Weight, fsnode_directory.id().clone()))
        }
    }
}

#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug, Serialize, Deserialize)]
pub enum Entry<T, L> {
    Tree(T),
    Leaf(L),
}

impl<T, L> Entry<T, L> {
    pub fn into_tree(self) -> Option<T> {
        match self {
            Entry::Tree(tree) => Some(tree),
            _ => None,
        }
    }

    pub fn into_leaf(self) -> Option<L> {
        match self {
            Entry::Leaf(leaf) => Some(leaf),
            _ => None,
        }
    }

    pub fn map_leaf<L2>(self, m: impl FnOnce(L) -> L2) -> Entry<T, L2> {
        match self {
            Entry::Tree(tree) => Entry::Tree(tree),
            Entry::Leaf(leaf) => Entry::Leaf(m(leaf)),
        }
    }

    pub fn map_tree<T2>(self, m: impl FnOnce(T) -> T2) -> Entry<T2, L> {
        match self {
            Entry::Tree(tree) => Entry::Tree(m(tree)),
            Entry::Leaf(leaf) => Entry::Leaf(leaf),
        }
    }

    pub fn is_tree(&self) -> bool {
        match self {
            Entry::Tree(_) => true,
            _ => false,
        }
    }
}

#[async_trait]
impl<T, L> Loadable for Entry<T, L>
where
    T: Loadable + Sync,
    L: Loadable + Sync,
{
    type Value = Entry<T::Value, L::Value>;

    async fn load<'a, B: Blobstore>(
        &'a self,
        ctx: &'a CoreContext,
        blobstore: &'a B,
    ) -> Result<Self::Value, LoadableError> {
        Ok(match self {
            Entry::Tree(tree_id) => Entry::Tree(tree_id.load(ctx, blobstore).await?),
            Entry::Leaf(leaf_id) => Entry::Leaf(leaf_id.load(ctx, blobstore).await?),
        })
    }
}

#[async_trait]
impl<T, L> Storable for Entry<T, L>
where
    T: Storable + Send,
    L: Storable + Send,
{
    type Key = Entry<T::Key, L::Key>;

    async fn store<'a, B: Blobstore>(
        self,
        ctx: &'a CoreContext,
        blobstore: &'a B,
    ) -> Result<Self::Key> {
        Ok(match self {
            Entry::Tree(tree) => Entry::Tree(tree.store(ctx, blobstore).await?),
            Entry::Leaf(leaf) => Entry::Leaf(leaf.store(ctx, blobstore).await?),
        })
    }
}

#[derive(Clone)]
pub struct PathTree<V> {
    pub value: V,
    pub subentries: BTreeMap<MPathElement, Self>,
}

impl<V> PathTree<V>
where
    V: Default,
{
    pub fn insert(&mut self, path: Option<MPath>, value: V) {
        let node = path.into_iter().flatten().fold(self, |node, element| {
            node.subentries
                .entry(element)
                .or_insert_with(Default::default)
        });
        node.value = value;
    }

    pub fn insert_and_merge<T>(&mut self, path: Option<MPath>, value: T)
    where
        V: Extend<T>,
    {
        let node = path.into_iter().flatten().fold(self, |node, element| {
            node.subentries
                .entry(element)
                .or_insert_with(Default::default)
        });
        node.value.extend(std::iter::once(value));
    }

    pub fn get(&self, path: Option<&MPath>) -> Option<&V> {
        let mut tree = self;
        for elem in path.into_iter().flatten() {
            match tree.subentries.get(elem) {
                Some(subtree) => tree = subtree,
                None => return None,
            }
        }
        Some(&tree.value)
    }

    pub fn insert_and_prune(&mut self, path: Option<MPath>, value: V) {
        let node = path.into_iter().flatten().fold(self, |node, element| {
            node.subentries
                .entry(element)
                .or_insert_with(Default::default)
        });
        node.value = value;
        node.subentries.clear();
    }
}

impl<V> Default for PathTree<V>
where
    V: Default,
{
    fn default() -> Self {
        Self {
            value: Default::default(),
            subentries: Default::default(),
        }
    }
}

impl<V> FromIterator<(MPath, V)> for PathTree<V>
where
    V: Default,
{
    fn from_iter<I>(iter: I) -> Self
    where
        I: IntoIterator<Item = (MPath, V)>,
    {
        let mut tree: Self = Default::default();
        for (path, value) in iter {
            tree.insert(Some(path), value);
        }
        tree
    }
}

impl<V> FromIterator<(Option<MPath>, V)> for PathTree<V>
where
    V: Default,
{
    fn from_iter<I>(iter: I) -> Self
    where
        I: IntoIterator<Item = (Option<MPath>, V)>,
    {
        let mut tree: Self = Default::default();
        for (path, value) in iter {
            tree.insert(path, value);
        }
        tree
    }
}

pub struct PathTreeIter<V> {
    frames: Vec<(Option<MPath>, PathTree<V>)>,
}

impl<V> Iterator for PathTreeIter<V> {
    type Item = (Option<MPath>, V);

    fn next(&mut self) -> Option<Self::Item> {
        let (path, PathTree { value, subentries }) = self.frames.pop()?;
        for (name, subentry) in subentries {
            self.frames.push((
                Some(MPath::join_opt_element(path.as_ref(), &name)),
                subentry,
            ));
        }
        Some((path, value))
    }
}

impl<V> IntoIterator for PathTree<V> {
    type Item = (Option<MPath>, V);
    type IntoIter = PathTreeIter<V>;

    fn into_iter(self) -> Self::IntoIter {
        PathTreeIter {
            frames: vec![(None, self)],
        }
    }
}

/// Traced allows you to trace a given parent through manifest derivation. For example, if you
/// assign ID 1 to a tree, then perform manifest derivation, then further entries you presented to
/// you that came from this parent will have the same ID.
#[derive(Debug)]
pub struct Traced<I, E>(Option<I>, E);

impl<I, E: Hash> Hash for Traced<I, E> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.1.hash(state);
    }
}

impl<I, E: PartialEq> PartialEq for Traced<I, E> {
    fn eq(&self, other: &Self) -> bool {
        self.1 == other.1
    }
}

impl<I, E: Eq> Eq for Traced<I, E> {}

impl<I: Copy, E: Copy> Copy for Traced<I, E> {}

impl<I: Clone, E: Clone> Clone for Traced<I, E> {
    fn clone(&self) -> Self {
        Self(self.0.clone(), self.1.clone())
    }
}

impl<I, E> Traced<I, E> {
    pub fn generate(e: E) -> Self {
        Self(None, e)
    }

    pub fn assign(i: I, e: E) -> Self {
        Self(Some(i), e)
    }

    pub fn id(&self) -> Option<&I> {
        self.0.as_ref()
    }

    pub fn untraced(&self) -> &E {
        &self.1
    }

    pub fn into_untraced(self) -> E {
        self.1
    }
}

impl<I: Copy, E> Traced<I, E> {
    fn inherit_into_entry<TreeId, LeafId>(
        &self,
        e: Entry<TreeId, LeafId>,
    ) -> Entry<Traced<I, TreeId>, Traced<I, LeafId>> {
        match e {
            Entry::Tree(t) => Entry::Tree(Traced(self.0, t)),
            Entry::Leaf(l) => Entry::Leaf(Traced(self.0, l)),
        }
    }
}

impl<I, TreeId, LeafId> From<Entry<Traced<I, TreeId>, Traced<I, LeafId>>>
    for Entry<TreeId, LeafId>
{
    fn from(entry: Entry<Traced<I, TreeId>, Traced<I, LeafId>>) -> Self {
        match entry {
            Entry::Tree(Traced(_, t)) => Entry::Tree(t),
            Entry::Leaf(Traced(_, l)) => Entry::Leaf(l),
        }
    }
}

impl<I: Send + Sync + Copy + 'static, M: Manifest> Manifest for Traced<I, M> {
    type TreeId = Traced<I, <M as Manifest>::TreeId>;
    type LeafId = Traced<I, <M as Manifest>::LeafId>;

    fn list(&self) -> Box<dyn Iterator<Item = (MPathElement, Entry<Self::TreeId, Self::LeafId>)>> {
        Box::new(
            self.1
                .list()
                .map(|(path, entry)| (path, self.inherit_into_entry(entry)))
                .collect::<Vec<_>>()
                .into_iter(),
        )
    }

    fn lookup(&self, name: &MPathElement) -> Option<Entry<Self::TreeId, Self::LeafId>> {
        self.1.lookup(name).map(|e| self.inherit_into_entry(e))
    }
}

#[async_trait]
impl<I: Clone + 'static + Send + Sync, M: Loadable + Send + Sync> Loadable for Traced<I, M> {
    type Value = Traced<I, <M as Loadable>::Value>;

    async fn load<'a, B: Blobstore>(
        &'a self,
        ctx: &'a CoreContext,
        blobstore: &'a B,
    ) -> Result<Self::Value, LoadableError> {
        let id = self.0.clone();
        let v = self.1.load(ctx, blobstore).await?;
        Ok(Traced(id, v))
    }
}


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

pub mod types;

mod queue;
pub use queue::AsyncMethodRequestQueue;
pub use queue::ClaimedBy;
pub use queue::RequestId;

pub mod tokens {
    pub use crate::types::MegarepoAddBranchingTargetToken;
    pub use crate::types::MegarepoAddTargetToken;
    pub use crate::types::MegarepoChangeTargetConfigToken;
    pub use crate::types::MegarepoRemergeSourceToken;
    pub use crate::types::MegarepoSyncChangesetToken;
}


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

use anyhow::anyhow;
use anyhow::Error;
use blobstore::Blobstore;
use blobstore::PutBehaviour;
use blobstore::Storable;
use bookmarks::BookmarkKey;
use context::CoreContext;
use futures::stream;
use futures::StreamExt;
use futures::TryStreamExt;
use megarepo_error::MegarepoError;
use memblob::Memblob;
use mononoke_api::Mononoke;
use mononoke_types::BlobstoreKey as BlobstoreKeyTrait;
use mononoke_types::RepositoryId;
use mononoke_types::Timestamp;
use requests_table::BlobstoreKey;
pub use requests_table::ClaimedBy;
use requests_table::LongRunningRequestEntry;
use requests_table::LongRunningRequestsQueue;
pub use requests_table::RequestId;
use requests_table::RequestType;
pub use requests_table::RowId;
use requests_table::SqlLongRunningRequestsQueue;
use sql_construct::SqlConstruct;

use crate::types::IntoApiFormat;
use crate::types::IntoConfigFormat;
use crate::types::MegarepoAsynchronousRequestParams;
use crate::types::MegarepoAsynchronousRequestResult;
use crate::types::Request;
use crate::types::ThriftParams;
use crate::types::Token;

const INITIAL_POLL_DELAY_MS: u64 = 1000;
const MAX_POLL_DURATION: Duration = Duration::from_secs(60);

#[derive(Clone)]
pub struct AsyncMethodRequestQueue {
    blobstore: Arc<dyn Blobstore>,
    table: Arc<dyn LongRunningRequestsQueue>,
}

impl AsyncMethodRequestQueue {
    pub fn new(table: Arc<dyn LongRunningRequestsQueue>, blobstore: Arc<dyn Blobstore>) -> Self {
        Self { blobstore, table }
    }

    pub fn new_test_in_memory() -> Result<Self, Error> {
        let blobstore: Arc<dyn Blobstore> = Arc::new(Memblob::new(PutBehaviour::IfAbsent));
        let table: Arc<dyn LongRunningRequestsQueue> =
            Arc::new(SqlLongRunningRequestsQueue::with_sqlite_in_memory()?);

        Ok(Self { blobstore, table })
    }

    pub async fn enqueue<P: ThriftParams>(
        &self,
        ctx: CoreContext,
        mononoke: &Mononoke,
        thrift_params: P,
    ) -> Result<<P::R as Request>::Token, Error> {
        let request_type = RequestType(P::R::NAME.to_owned());
        let target = thrift_params
            .target()
            .clone()
            .into_config_format(mononoke)?;
        let rust_params: MegarepoAsynchronousRequestParams = thrift_params.into();
        let params_object_id = rust_params.store(&ctx, &self.blobstore).await?;
        let blobstore_key = BlobstoreKey(params_object_id.blobstore_key());
        let table_id = self
            .table
            .add_request(
                &ctx,
                &request_type,
                &RepositoryId::new(i32::try_from(target.repo_id)?),
                &BookmarkKey::new(&target.bookmark)?,
                &blobstore_key,
            )
            .await?;
        let token = <P::R as Request>::Token::from_db_id_and_target(
            table_id,
            target.into_api_format(mononoke)?,
        );
        Ok(token)
    }

    pub async fn dequeue(
        &self,
        ctx: &CoreContext,
        claimed_by: &ClaimedBy,
        supported_repos: &[RepositoryId],
    ) -> Result<Option<(RequestId, MegarepoAsynchronousRequestParams)>, MegarepoError> {
        let entry = self
            .table
            .claim_and_get_new_request(ctx, claimed_by, supported_repos)
            .await?;

        if let Some(entry) = entry {
            let thrift_params = MegarepoAsynchronousRequestParams::load_from_key(
                ctx,
                &self.blobstore,
                &entry.args_blobstore_key.0,
            )
            .await?;
            let req_id = RequestId(entry.id, entry.request_type);
            Ok(Some((req_id, thrift_params)))
        } else {
            // empty queue
            Ok(None)
        }
    }

    pub async fn complete(
        &self,
        ctx: &CoreContext,
        req_id: &RequestId,
        result: MegarepoAsynchronousRequestResult,
    ) -> Result<bool, MegarepoError> {
        let result_object_id = result.store(ctx, &self.blobstore).await?;
        let blobstore_key = BlobstoreKey(result_object_id.blobstore_key());
        Ok(self.table.mark_ready(ctx, req_id, blobstore_key).await?)
    }

    async fn poll_once<R: Request>(
        &self,
        ctx: &CoreContext,
        req_id: &RequestId,
    ) -> Result<Option<<R as Request>::ThriftResult>, MegarepoError> {
        let maybe_result_blobstore_key = match self.table.poll(ctx, req_id).await? {
            None => return Ok(None),
            Some((_, entry)) => entry.result_blobstore_key,
        };

        let result_blobstore_key = match maybe_result_blobstore_key {
            Some(rbk) => rbk,
            None => {
                return Err(MegarepoError::internal(anyhow!(
                    "Programming error: successful poll with empty result_blobstore_key for {:?}",
                    req_id
                )));
            }
        };

        let result: MegarepoAsynchronousRequestResult =
            MegarepoAsynchronousRequestResult::load_from_key(
                ctx,
                &self.blobstore,
                &result_blobstore_key.0,
            )
            .await?;
        Ok(Some(result.try_into()?))
    }

    pub async fn poll<T: Token>(
        &self,
        ctx: CoreContext,
        token: T,
    ) -> Result<<T::R as Request>::PollResponse, MegarepoError> {
        let mut backoff_ms = INITIAL_POLL_DELAY_MS;
        let before = Instant::now();
        let (row_id, _target) = token.to_db_id_and_target()?;
        let req_id = RequestId(row_id, RequestType(T::R::NAME.to_owned()));

        loop {
            let maybe_thrift_result: Option<<T::R as Request>::ThriftResult> =
                self.poll_once::<T::R>(&ctx, &req_id).await?;
            let next_sleep = Duration::from_millis(rand::random::<u64>() % backoff_ms);

            match maybe_thrift_result {
                Some(thrift_result) => {
                    // Nice, the result is ready!
                    return Ok(<T::R as Request>::thrift_result_into_poll_response(
                        thrift_result,
                    ));
                }
                None if before.elapsed() + next_sleep > MAX_POLL_DURATION => {
                    // The result is not yet ready, but we're out of time
                    return Ok(T::R::empty_poll_response());
                }
                None => {
                    // The result is not yet ready and we can wait a little longer
                    tokio::time::sleep(next_sleep).await;
                    backoff_ms *= 2;
                }
            }
        }
    }

    pub async fn update_in_progress_timestamp(
        &self,
        ctx: &CoreContext,
        req_id: &RequestId,
    ) -> Result<bool, MegarepoError> {
        Ok(self.table.update_in_progress_timestamp(ctx, req_id).await?)
    }

    pub async fn find_abandoned_requests(
        &self,
        ctx: &CoreContext,
        repo_ids: &[RepositoryId],
        abandoned_timestamp: Timestamp,
    ) -> Result<Vec<RequestId>, MegarepoError> {
        Ok(self
            .table
            .find_abandoned_requests(ctx, repo_ids, abandoned_timestamp)
            .await?)
    }

    pub async fn mark_abandoned_request_as_new(
        &self,
        ctx: &CoreContext,
        request_id: RequestId,
        abandoned_timestamp: Timestamp,
    ) -> Result<bool, MegarepoError> {
        Ok(self
            .table
            .mark_abandoned_request_as_new(ctx, request_id, abandoned_timestamp)
            .await?)
    }

    pub async fn requeue(
        &self,
        ctx: &CoreContext,
        request_id: RequestId,
    ) -> Result<bool, MegarepoError> {
        Ok(self.table.mark_new(ctx, &request_id).await?)
    }

    pub async fn list_requests(
        &self,
        ctx: &CoreContext,
        repo_ids: &[RepositoryId],
        last_update_newer_than: Option<&Timestamp>,
    ) -> Result<
        Vec<(
            RequestId,
            LongRunningRequestEntry,
            MegarepoAsynchronousRequestParams,
        )>,
        MegarepoError,
    > {
        let entries = self
            .table
            .list_requests(ctx, repo_ids, last_update_newer_than)
            .await?;

        stream::iter(entries)
            .map(|entry| async {
                let thrift_params = MegarepoAsynchronousRequestParams::load_from_key(
                    ctx,
                    &self.blobstore,
                    &entry.args_blobstore_key.0,
                )
                .await?;
                let req_id = RequestId(entry.id.clone(), entry.request_type.clone());
                Ok::<_, MegarepoError>((req_id, entry, thrift_params))
            })
            .buffer_unordered(10)
            .try_collect()
            .await
    }

    pub async fn get_request_by_id(
        &self,
        ctx: &CoreContext,
        row_id: &RowId,
    ) -> Result<
        Option<(
            RequestId,
            LongRunningRequestEntry,
            MegarepoAsynchronousRequestParams,
            Option<MegarepoAsynchronousRequestResult>,
        )>,
        MegarepoError,
    > {
        let entry = self.table.test_get_request_entry_by_id(ctx, row_id).await?;

        if let Some(entry) = entry {
            let thrift_params = MegarepoAsynchronousRequestParams::load_from_key(
                ctx,
                &self.blobstore,
                &entry.args_blobstore_key.0,
            )
            .await?;
            let req_id = RequestId(entry.id.clone(), entry.request_type.clone());
            let thrift_result = if let Some(result_blobstore_key) = &entry.result_blobstore_key {
                Some(
                    MegarepoAsynchronousRequestResult::load_from_key(
                        ctx,
                        &self.blobstore,
                        &result_blobstore_key.0,
                    )
                    .await?,
                )
            } else {
                None
            };
            Ok(Some((req_id, entry, thrift_params, thrift_result)))
        } else {
            // empty queue
            Ok(None)
        }
    }
}

#[cfg(test)]
mod tests {
    use context::CoreContext;
    use fbinit::FacebookInit;
    use mononoke_api::Repo;
    use requests_table::ClaimedBy;
    use requests_table::RequestStatus;
    use source_control::types::MegarepoTarget as ThriftMegarepoTarget;
    use source_control::types::RepoSpecifier;
    use source_control::MegarepoAddBranchingTargetParams as ThriftMegarepoAddBranchingTargetParams;
    use source_control::MegarepoAddBranchingTargetResult;
    use source_control::MegarepoAddTargetParams as ThriftMegarepoAddTargetParams;
    use source_control::MegarepoAddTargetResult;
    use source_control::MegarepoChangeTargetConfigParams as ThriftMegarepoChangeTargetConfigParams;
    use source_control::MegarepoChangeTargetConfigResult;
    use source_control::MegarepoRemergeSourceParams as ThriftMegarepoRemergeSourceParams;
    use source_control::MegarepoRemergeSourceResult;
    use source_control::MegarepoSyncChangesetParams as ThriftMegarepoSyncChangesetParams;
    use source_control::MegarepoSyncChangesetResult;
    use source_control::MegarepoSyncTargetConfig as ThriftMegarepoSyncTargetConfig;

    use super::*;
    use crate::types::MegarepoAddBranchingSyncTarget;
    use crate::types::MegarepoAddSyncTarget;
    use crate::types::MegarepoAsynchronousRequestResult;
    use crate::types::MegarepoChangeTargetConfig;
    use crate::types::MegarepoRemergeSource;
    use crate::types::MegarepoSyncChangeset;

    macro_rules! test_enqueue_dequeue_and_poll_once {
        {
            $fn_name: ident,
            $request_struct: ident,
            $thrift_params: expr,
            $result: ident,
            $request_type: expr,
        } => {
            #[fbinit::test]
            async fn $fn_name(fb: FacebookInit) -> Result<(), Error> {
                let q = AsyncMethodRequestQueue::new_test_in_memory().unwrap();
                let ctx = CoreContext::test_mock(fb);
                let repo: Repo = test_repo_factory::build_empty(ctx.fb).await?;
                let mononoke =
                    Mononoke::new_test(vec![("test".to_string(), repo)]).await?;

                // Enqueue a request
                let params = $thrift_params;
                let token = q.enqueue(ctx.clone(), &mononoke, params.clone()).await?;

                // Verify that request metadata is in the db and has expected values
                let (row_id, _) = token.to_db_id_and_target()?;
                let entry = q
                    .table
                    .test_get_request_entry_by_id(&ctx, &row_id)
                    .await?
                    .expect("Request is mising in the DB");
                assert_eq!(entry.status, RequestStatus::New);
                assert_eq!(entry.started_processing_at, None);
                assert_eq!(entry.ready_at, None);
                assert_eq!(entry.polled_at, None);
                assert_eq!(entry.repo_id,  RepositoryId::new(0));
                assert_eq!(
                    entry.request_type,
                    RequestType($request_type.to_string())
                );
                let req_id = RequestId(row_id, entry.request_type);

                // Verify that poll_once on this request in a "new" state
                // returns None
                let new_poll = q.poll_once::<$request_struct>(&ctx, &req_id).await?;
                assert!(new_poll.is_none());

                // Simulate the tailer and grab the element from the queue, this should return the params
                // back and flip its state back to "in_progress"
                let (req_id, params_from_store) = q.dequeue(&ctx, &ClaimedBy("tests".to_string()), &[entry.repo_id]).await?.unwrap();

                // Verify that request params from blobstore match what we put there
                assert_eq!(params_from_store, params.into());

                // Verify that request params are in the blobstore

                // Verify that poll_once on this request in a "in_progress" state
                // returns None
                let in_progress_poll = q.poll_once::<$request_struct>(&ctx,  &req_id).await?;
                assert!(in_progress_poll.is_none());

                // Inject a result for this request
                // Verify that poll_once on this request in a "in_progress" state
                // returns injected result
                let fake_specific_result: $result = Default::default();
                let fake_result: MegarepoAsynchronousRequestResult = fake_specific_result.clone().into();
                q.complete(&ctx, &req_id, fake_result).await?;
                let ready_poll = q.poll_once::<$request_struct>(&ctx, &req_id).await?;
                let ready_poll_response = ready_poll.unwrap();
                assert_eq!(ready_poll_response, fake_specific_result);

                // After a successful poll, request is marked as polled
                let entry = q.table.test_get_request_entry_by_id(&ctx, &req_id.0).await?.unwrap();
                assert_eq!(entry.status, RequestStatus::Polled);

                Ok(())
            }
        }
    }

    test_enqueue_dequeue_and_poll_once! {
        test_enqueue_dequeue_and_poll_once_add_target,
        MegarepoAddSyncTarget,
        ThriftMegarepoAddTargetParams {
            config_with_new_target: ThriftMegarepoSyncTargetConfig {
                target: ThriftMegarepoTarget {
                    bookmark: "oculus".to_string(),
                    repo_id: Some(0),
                    ..Default::default()
                },
                ..Default::default()
            },
            ..Default::default()
        },
        MegarepoAddTargetResult,
        "megarepo_add_sync_target",
    }

    test_enqueue_dequeue_and_poll_once! {
        test_enqueue_dequeue_and_poll_once_add_branching_target,
        MegarepoAddBranchingSyncTarget,
        ThriftMegarepoAddBranchingTargetParams {
            target: ThriftMegarepoTarget {
                bookmark: "oculus".to_string(),
                repo_id: Some(0),
                ..Default::default()
            },
            ..Default::default()
        },
        MegarepoAddBranchingTargetResult,
        "megarepo_add_branching_sync_target",
    }

    test_enqueue_dequeue_and_poll_once! {
        test_enqueue_dequeue_and_poll_once_sync_changeset,
        MegarepoSyncChangeset,
        ThriftMegarepoSyncChangesetParams {
            target: ThriftMegarepoTarget {
                bookmark: "oculus".to_string(),
                repo: Some(
                    RepoSpecifier {
                        name: "test".to_string(),
                        ..Default::default()
                    }
                ),
                ..Default::default()
            },
            ..Default::default()
        },
        MegarepoSyncChangesetResult,
        "megarepo_sync_changeset",
    }

    test_enqueue_dequeue_and_poll_once! {
        test_enqueue_dequeue_and_poll_once_change_config,
        MegarepoChangeTargetConfig,
        ThriftMegarepoChangeTargetConfigParams {
            target: ThriftMegarepoTarget {
                bookmark: "oculus".to_string(),
                repo_id: Some(0),
                ..Default::default()
            },
            ..Default::default()
        },
        MegarepoChangeTargetConfigResult,
        "megarepo_change_target_config",
    }

    test_enqueue_dequeue_and_poll_once! {
        test_enqueue_dequeue_and_poll_once_remerge_source,
        MegarepoRemergeSource,
        ThriftMegarepoRemergeSourceParams {
            target: ThriftMegarepoTarget {
                bookmark: "oculus".to_string(),
                repo_id: Some(0),
                ..Default::default()
            },
            ..Default::default()
        },
        MegarepoRemergeSourceResult,
        "megarepo_remerge_source",
    }
}


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

use anyhow::anyhow;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::impl_loadable_storable;
use blobstore::Blobstore;
use context::CoreContext;
use fbthrift::compact_protocol;
pub use megarepo_config::SyncTargetConfig;
pub use megarepo_config::Target;
use megarepo_error::MegarepoError;
pub use megarepo_types_thrift::MegarepoAsynchronousRequestParams as ThriftMegarepoAsynchronousRequestParams;
pub use megarepo_types_thrift::MegarepoAsynchronousRequestParamsId as ThriftMegarepoAsynchronousRequestParamsId;
pub use megarepo_types_thrift::MegarepoAsynchronousRequestResult as ThriftMegarepoAsynchronousRequestResult;
pub use megarepo_types_thrift::MegarepoAsynchronousRequestResultId as ThriftMegarepoAsynchronousRequestResultId;
use mononoke_api::Mononoke;
use mononoke_types::hash::Blake2;
use mononoke_types::impl_typed_context;
use mononoke_types::impl_typed_hash_no_context;
use mononoke_types::BlobstoreKey;
use mononoke_types::RepositoryId;
pub use requests_table::RequestStatus;
pub use requests_table::RequestType;
pub use requests_table::RowId;
pub use source_control::MegarepoAddBranchingTargetParams as ThriftMegarepoAddBranchingTargetParams;
pub use source_control::MegarepoAddBranchingTargetPollResponse as ThriftMegarepoAddBranchingTargetPollResponse;
pub use source_control::MegarepoAddBranchingTargetResponse as ThriftMegarepoAddBranchingTargetResponse;
pub use source_control::MegarepoAddBranchingTargetResult as ThriftMegarepoAddBranchingTargetResult;
pub use source_control::MegarepoAddBranchingTargetToken as ThriftMegarepoAddBranchingTargetToken;
pub use source_control::MegarepoAddTargetParams as ThriftMegarepoAddTargetParams;
pub use source_control::MegarepoAddTargetPollResponse as ThriftMegarepoAddTargetPollResponse;
pub use source_control::MegarepoAddTargetResponse as ThriftMegarepoAddTargetResponse;
pub use source_control::MegarepoAddTargetResult as ThriftMegarepoAddTargetResult;
pub use source_control::MegarepoAddTargetToken as ThriftMegarepoAddTargetToken;
pub use source_control::MegarepoChangeConfigToken as ThriftMegarepoChangeConfigToken;
pub use source_control::MegarepoChangeTargetConfigParams as ThriftMegarepoChangeTargetConfigParams;
pub use source_control::MegarepoChangeTargetConfigPollResponse as ThriftMegarepoChangeTargetConfigPollResponse;
pub use source_control::MegarepoChangeTargetConfigResponse as ThriftMegarepoChangeTargetConfigResponse;
pub use source_control::MegarepoChangeTargetConfigResult as ThriftMegarepoChangeTargetConfigResult;
pub use source_control::MegarepoRemergeSourceParams as ThriftMegarepoRemergeSourceParams;
pub use source_control::MegarepoRemergeSourcePollResponse as ThriftMegarepoRemergeSourcePollResponse;
pub use source_control::MegarepoRemergeSourceResponse as ThriftMegarepoRemergeSourceResponse;
pub use source_control::MegarepoRemergeSourceResult as ThriftMegarepoRemergeSourceResult;
pub use source_control::MegarepoRemergeSourceToken as ThriftMegarepoRemergeSourceToken;
pub use source_control::MegarepoSyncChangesetParams as ThriftMegarepoSyncChangesetParams;
pub use source_control::MegarepoSyncChangesetPollResponse as ThriftMegarepoSyncChangesetPollResponse;
pub use source_control::MegarepoSyncChangesetResponse as ThriftMegarepoSyncChangesetResponse;
pub use source_control::MegarepoSyncChangesetResult as ThriftMegarepoSyncChangesetResult;
pub use source_control::MegarepoSyncChangesetToken as ThriftMegarepoSyncChangesetToken;
pub use source_control::MegarepoSyncTargetConfig as ThriftMegarepoSyncTargetConfig;
pub use source_control::MegarepoTarget as ThriftMegarepoTarget;
pub use source_control::RepoSpecifier as ThriftRepoSpecifier;

/// Grouping of types and behaviors for an asynchronous request
pub trait Request: Sized + Send + Sync {
    /// Name of the request
    const NAME: &'static str;
    /// Rust newtype for a polling token
    type Token: Token;

    /// Underlying thrift type for request params
    type ThriftParams: ThriftParams<R = Self>;

    /// Underlying thrift type for successful request response
    type ThriftResponse;

    /// Underlying thrift type for for request result (response or error)
    type ThriftResult: ThriftResult<R = Self>;

    /// A type representing potentially present response
    type PollResponse;

    /// Convert thrift result into a result of a poll response
    fn thrift_result_into_poll_response(tr: Self::ThriftResult) -> Self::PollResponse;

    /// Return an empty poll response. This indicates
    /// that the request hasn't been processed yet
    fn empty_poll_response() -> Self::PollResponse;
}

/// Thrift type representing async service method parameters
pub trait ThriftParams: Sized + Send + Sync + Into<MegarepoAsynchronousRequestParams> {
    type R: Request<ThriftParams = Self>;

    /// Every *Params argument referes to some Target
    /// This method is needed to extract it from the
    /// implementor of this trait
    fn target(&self) -> &ThriftMegarepoTarget;
}
pub trait ThriftResult:
    Sized + Send + Sync + TryFrom<MegarepoAsynchronousRequestResult, Error = MegarepoError>
{
    type R: Request<ThriftResult = Self>;
}

/// Polling token for an async service method
pub trait Token: Clone + Sized + Send + Sync {
    type R: Request<Token = Self>;
    type ThriftToken;

    fn into_thrift(self) -> Self::ThriftToken;
    fn from_db_id_and_target(id: RowId, target: ThriftMegarepoTarget) -> Self;
    fn to_db_id_and_target(&self) -> Result<(RowId, ThriftMegarepoTarget), MegarepoError>;

    /// Every Token referes to some Target
    /// This method is needed to extract it from the
    /// implementor of this trait
    fn target(&self) -> &ThriftMegarepoTarget;
}

/// This macro implements an async service method type,
/// which can be stored/retrieved from the blobstore.
/// Such types are usually represented as value/handle pairs.
/// Since we need to implement (potentially foreign) traits
/// on these types, we also define corrensponding Rust types
/// Some of the defined types (like context or thrift_type_newtype)
/// are not used from outside of the macro, but we still need
/// to pass identifiers for them from the outside, because
/// Rusts' macro hygiene does not allow identifier generation Â¯\_(ãƒ„)_/Â¯
macro_rules! impl_async_svc_stored_type {
    {
        /// Rust type for the Loadable handle
        handle_type => $handle_type: ident,
        /// Underlying thrift type for the handle
        handle_thrift_type => $handle_thrift_type: ident,
        /// A name for a Newtype-style trait, required by `impl_typed_hash_no_context`
        /// Rust type for the Storable value
        value_type => $value_type: ident,
        /// Underlying thrift type for the value
        value_thrift_type => $value_thrift_type: ident,
        /// A helper struct for hash computations
        context_type => $context_type: ident,
    } => {
        /// Rust handle type, wrapper around a Blake2 instance
        #[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Debug, Hash)]
        pub struct $handle_type(Blake2);

        impl_typed_hash_no_context! {
            hash_type => $handle_type,
            thrift_type => $handle_thrift_type,
            blobstore_key => concat!("async.svc.", stringify!($value_type)),
        }

        // Typed context type is needed for hash computation
        impl_typed_context! {
            hash_type => $handle_type,
            context_type => $context_type,
            context_key => stringify!($value_type),
        }

        /// Main value type
        #[derive(Debug, Clone, PartialEq)]
        pub struct $value_type {
            id: $handle_type,
            thrift: $value_thrift_type,
        }

        impl $value_type {
            pub fn from_thrift(thrift: $value_thrift_type) -> Self {
                let data = compact_protocol::serialize(&thrift);
                let mut context = $context_type::new();
                context.update(&data);
                let id = context.finish();
                Self { id, thrift }
            }

            pub async fn load_from_key(ctx: &CoreContext, blobstore: &Arc<dyn Blobstore>, key: &str) -> Result<Self, MegarepoError> {
                let bytes = blobstore.get(ctx, key).await?;

                let prefix = concat!("async.svc.", stringify!($value_type), ".blake2.");
                if key.strip_prefix(prefix).is_none() {
                    return Err(MegarepoError::internal(anyhow!("{} is not a blobstore key for {}", key, stringify!($value_type))));
                }

                match bytes {
                    Some(bytes) => Ok(bytes.into_bytes().try_into()?),
                    None => Err(MegarepoError::internal(anyhow!("Missing blob: {}", key))),
                }
            }

            pub fn handle(&self) -> &$handle_type {
                &self.id
            }

            pub fn thrift(&self) -> &$value_thrift_type {
                &self.thrift
            }

        }

        // Conversions between thrift types and their Rust counterparts

        impl TryFrom<$handle_thrift_type> for $handle_type {
            type Error = Error;

            fn try_from(t: $handle_thrift_type) -> Result<Self, Self::Error> {
                Self::from_thrift(t)
            }
        }

        impl From<$handle_type> for $handle_thrift_type {
            fn from(other: $handle_type) -> Self {
                Self(mononoke_types_thrift::IdType::Blake2(other.0.into_thrift()))
            }
        }

        impl TryFrom<$value_thrift_type> for $value_type {
            type Error = Error;

            fn try_from(t: $value_thrift_type) -> Result<Self, Self::Error> {
                Ok(Self::from_thrift(t))
            }
        }

        impl From<$value_type> for $value_thrift_type {
            fn from(other: $value_type) -> Self {
                other.thrift
            }
        }

        impl_loadable_storable! {
            handle_type => $handle_type,
            handle_thrift_type => $handle_thrift_type,
            value_type => $value_type,
            value_thrift_type => $value_thrift_type,
        }
    }
}

/// A macro to call impl_async_svc_stored_type for params/result
/// types, as well as define a bunch of relationships between
/// these types, and their Request-related frients.
/// An underlying idea is to define as much behavior and relationships
/// as possible in the type system, so that we
/// (a) minimize a chance of using incorrent pair of types somewhere
/// (b) can write generic enqueing/polling functions
macro_rules! impl_async_svc_method_types {
    {
        method_name => $method_name: expr,
        request_struct => $request_struct: ident,

        params_value_thrift_type => $params_value_thrift_type: ident,
        params_union_variant => $params_union_variant: ident,

        result_value_thrift_type => $result_value_thrift_type: ident,
        result_union_variant => $result_union_variant: ident,

        response_type => $response_type: ident,
        poll_response_type => $poll_response_type: ident,
        token_type => $token_type: ident,
        token_thrift_type => $token_thrift_type: ident,

        fn target(&$self_ident: ident: ThriftParams) -> &ThriftMegarepoTarget $target_in_params: tt

    } => {
        impl ThriftParams for $params_value_thrift_type {
            type R = $request_struct;

            fn target(&$self_ident) -> &ThriftMegarepoTarget {
                $target_in_params
            }
        }

        #[derive(Clone)]
        pub struct $token_type(pub $token_thrift_type);

        impl Token for $token_type {
            type ThriftToken = $token_thrift_type;
            type R = $request_struct;

            fn from_db_id_and_target(id: RowId, target: ThriftMegarepoTarget) -> Self {
                // Thrift token is a string alias
                // but's guard ourselves here against
                // it changing unexpectedly.
                let thrift_token = $token_thrift_type {
                    target,
                    id: id.0 as i64,
                    ..Default::default()
                };
                Self(thrift_token)
            }

            fn to_db_id_and_target(&self) -> Result<(RowId, ThriftMegarepoTarget), MegarepoError> {
                let row_id = self.0.id as u64;
                let row_id = RowId(row_id);
                let target = self.0.target.clone();

                Ok((row_id, target))
            }

            fn into_thrift(self) -> $token_thrift_type {
                self.0
            }

            fn target(&self) -> &ThriftMegarepoTarget {
                &self.0.target
            }
        }

        impl From<Result<$response_type, MegarepoError>> for MegarepoAsynchronousRequestResult {
            fn from(r: Result<$response_type, MegarepoError>) -> MegarepoAsynchronousRequestResult {
                let thrift = match r {
                    Ok(payload) => ThriftMegarepoAsynchronousRequestResult::$result_union_variant($result_value_thrift_type::success(payload)),
                    Err(e) => ThriftMegarepoAsynchronousRequestResult::$result_union_variant($result_value_thrift_type::error(e.into()))
                };

                MegarepoAsynchronousRequestResult::from_thrift(thrift)
            }
        }

        impl From<$result_value_thrift_type> for MegarepoAsynchronousRequestResult {
            fn from(r: $result_value_thrift_type) -> MegarepoAsynchronousRequestResult {
                let thrift = ThriftMegarepoAsynchronousRequestResult::$result_union_variant(r);
                MegarepoAsynchronousRequestResult::from_thrift(thrift)
            }
        }

        impl From<$params_value_thrift_type> for MegarepoAsynchronousRequestParams{
            fn from(params: $params_value_thrift_type) -> MegarepoAsynchronousRequestParams {
                MegarepoAsynchronousRequestParams::from_thrift(
                    ThriftMegarepoAsynchronousRequestParams::$params_union_variant(params)
                )
            }
        }

        impl ThriftResult for $result_value_thrift_type {
            type R = $request_struct;
        }

        impl TryFrom<MegarepoAsynchronousRequestResult> for $result_value_thrift_type {
            type Error = MegarepoError;

            fn try_from(r: MegarepoAsynchronousRequestResult) -> Result<$result_value_thrift_type, Self::Error> {
                match r.thrift {
                    ThriftMegarepoAsynchronousRequestResult::$result_union_variant(payload) => Ok(payload),
                    ThriftMegarepoAsynchronousRequestResult::UnknownField(x) => {
                        // TODO: maybe use structured error?
                        Err(MegarepoError::internal(
                            anyhow!(
                                "failed to parse {} thrift. UnknownField: {}",
                                stringify!($result_value_thrift_type),
                                x,
                            )
                        ))
                    },
                    x => {
                        Err(MegarepoError::internal(
                            anyhow!(
                                "failed to parse {} thrift. The result union contains the wrong result variant: {:?}",
                                stringify!($result_value_thrift_type),
                                x,
                            )
                        ))
                    }
                }
            }
        }

        pub struct $request_struct;

        impl Request for $request_struct {
            const NAME: &'static str = $method_name;

            type Token = $token_type;
            type ThriftParams = $params_value_thrift_type;
            type ThriftResult = $result_value_thrift_type;
            type ThriftResponse = $response_type;

            type PollResponse = $poll_response_type;

            fn thrift_result_into_poll_response(
                thrift_result: Self::ThriftResult,
            ) -> Self::PollResponse {
                $poll_response_type { result: Some(thrift_result), ..Default::default() }
            }

            fn empty_poll_response() -> Self::PollResponse {
                $poll_response_type { result: None, ..Default::default() }
            }
        }

    }
}

// Params and result types for megarepo_add_sync_target

impl_async_svc_method_types! {
    method_name => "megarepo_add_sync_target",
    request_struct => MegarepoAddSyncTarget,

    params_value_thrift_type => ThriftMegarepoAddTargetParams,
    params_union_variant => megarepo_add_target_params,

    result_value_thrift_type => ThriftMegarepoAddTargetResult,
    result_union_variant => megarepo_add_target_result,

    response_type => ThriftMegarepoAddTargetResponse,
    poll_response_type => ThriftMegarepoAddTargetPollResponse,
    token_type => MegarepoAddTargetToken,
    token_thrift_type => ThriftMegarepoAddTargetToken,

    fn target(&self: ThriftParams) -> &ThriftMegarepoTarget {
        &self.config_with_new_target.target
    }
}

// Params and result types for megarepo_add_branching_sync_target

impl_async_svc_method_types! {
    method_name => "megarepo_add_branching_sync_target",
    request_struct => MegarepoAddBranchingSyncTarget,

    params_value_thrift_type => ThriftMegarepoAddBranchingTargetParams,
    params_union_variant => megarepo_add_branching_target_params,

    result_value_thrift_type => ThriftMegarepoAddBranchingTargetResult,
    result_union_variant => megarepo_add_branching_target_result,

    response_type => ThriftMegarepoAddBranchingTargetResponse,
    poll_response_type => ThriftMegarepoAddBranchingTargetPollResponse,
    token_type => MegarepoAddBranchingTargetToken,
    token_thrift_type => ThriftMegarepoAddBranchingTargetToken,

    fn target(&self: ThriftParams) -> &ThriftMegarepoTarget {
        &self.target
    }
}

// Params and result types for megarepo_change_target_config

impl_async_svc_method_types! {
    method_name => "megarepo_change_target_config",
    request_struct => MegarepoChangeTargetConfig,

    params_value_thrift_type => ThriftMegarepoChangeTargetConfigParams,
    params_union_variant => megarepo_change_target_params,

    result_value_thrift_type => ThriftMegarepoChangeTargetConfigResult,
    result_union_variant => megarepo_change_target_result,

    response_type => ThriftMegarepoChangeTargetConfigResponse,
    poll_response_type => ThriftMegarepoChangeTargetConfigPollResponse,
    token_type => MegarepoChangeTargetConfigToken,
    token_thrift_type => ThriftMegarepoChangeConfigToken,

    fn target(&self: ThriftParams) -> &ThriftMegarepoTarget {
        &self.target
    }
}

// Params and result types for megarepo_sync_changeset

impl_async_svc_method_types! {
    method_name => "megarepo_sync_changeset",
    request_struct => MegarepoSyncChangeset,

    params_value_thrift_type => ThriftMegarepoSyncChangesetParams,
    params_union_variant => megarepo_sync_changeset_params,

    result_value_thrift_type => ThriftMegarepoSyncChangesetResult,
    result_union_variant => megarepo_sync_changeset_result,

    response_type => ThriftMegarepoSyncChangesetResponse,
    poll_response_type => ThriftMegarepoSyncChangesetPollResponse,
    token_type => MegarepoSyncChangesetToken,
    token_thrift_type => ThriftMegarepoSyncChangesetToken,

    fn target(&self: ThriftParams) -> &ThriftMegarepoTarget {
        &self.target
    }
}

// Params and result types for megarepo_remerge_source

impl_async_svc_method_types! {
    method_name => "megarepo_remerge_source",
    request_struct => MegarepoRemergeSource,

    params_value_thrift_type => ThriftMegarepoRemergeSourceParams,
    params_union_variant => megarepo_remerge_source_params,

    result_value_thrift_type => ThriftMegarepoRemergeSourceResult,
    result_union_variant => megarepo_remerge_source_result,

    response_type => ThriftMegarepoRemergeSourceResponse,
    poll_response_type => ThriftMegarepoRemergeSourcePollResponse,
    token_type => MegarepoRemergeSourceToken,
    token_thrift_type => ThriftMegarepoRemergeSourceToken,

    fn target(&self: ThriftParams) -> &ThriftMegarepoTarget {
        &self.target
    }
}

impl_async_svc_stored_type! {
    handle_type => MegarepoAsynchronousRequestParamsId,
    handle_thrift_type => ThriftMegarepoAsynchronousRequestParamsId,
    value_type => MegarepoAsynchronousRequestParams,
    value_thrift_type => ThriftMegarepoAsynchronousRequestParams,
    context_type => MegarepoAsynchronousRequestParamsIdContext,
}

impl_async_svc_stored_type! {
    handle_type => MegarepoAsynchronousRequestResultId,
    handle_thrift_type => ThriftMegarepoAsynchronousRequestResultId,
    value_type => MegarepoAsynchronousRequestResult,
    value_thrift_type => ThriftMegarepoAsynchronousRequestResult,
    context_type => MegarepoAsynchronousRequestResultIdContext,
}

impl MegarepoAsynchronousRequestParams {
    pub fn target(&self) -> Result<&ThriftMegarepoTarget, MegarepoError> {
        match &self.thrift {
            ThriftMegarepoAsynchronousRequestParams::megarepo_add_target_params(params) => {
                Ok(params.target())
            }
            ThriftMegarepoAsynchronousRequestParams::megarepo_add_branching_target_params(
                params,
            ) => Ok(params.target()),
            ThriftMegarepoAsynchronousRequestParams::megarepo_change_target_params(params) => {
                Ok(params.target())
            }
            ThriftMegarepoAsynchronousRequestParams::megarepo_remerge_source_params(params) => {
                Ok(params.target())
            }
            ThriftMegarepoAsynchronousRequestParams::megarepo_sync_changeset_params(params) => {
                Ok(params.target())
            }
            ThriftMegarepoAsynchronousRequestParams::UnknownField(union_tag) => {
                Err(MegarepoError::internal(anyhow!(
                    "this type of reuqest (MegarepoAsynchronousRequestParams tag {}) not supported by this worker!",
                    union_tag
                )))
            }
        }
    }
}

/// Convert an item into a thrift type we use for storing configuration
pub trait IntoConfigFormat<T> {
    fn into_config_format(self, mononoke: &Mononoke) -> Result<T, MegarepoError>;
}

impl IntoConfigFormat<Target> for ThriftMegarepoTarget {
    fn into_config_format(self, mononoke: &Mononoke) -> Result<Target, MegarepoError> {
        let repo_id = match (self.repo, self.repo_id) {
            (Some(repo), _) => mononoke
                .repo_id_from_name(repo.name.clone())
                .ok_or_else(|| anyhow!("Invalid repo_name {}", repo.name))?
                .id() as i64,
            (_, Some(repo_id)) => repo_id,
            (None, None) => Err(anyhow!("both repo_id and repo_name are None!"))?,
        };

        Ok(Target {
            repo_id,
            bookmark: self.bookmark,
        })
    }
}

impl IntoConfigFormat<SyncTargetConfig> for ThriftMegarepoSyncTargetConfig {
    fn into_config_format(self, mononoke: &Mononoke) -> Result<SyncTargetConfig, MegarepoError> {
        Ok(SyncTargetConfig {
            target: self.target.into_config_format(mononoke)?,
            sources: self.sources,
            version: self.version,
        })
    }
}

/// Convert an item into a thrift type we use in APIs
pub trait IntoApiFormat<T> {
    fn into_api_format(self, mononoke: &Mononoke) -> Result<T, MegarepoError>;
}

#[async_trait]
impl IntoApiFormat<ThriftMegarepoTarget> for Target {
    fn into_api_format(self, mononoke: &Mononoke) -> Result<ThriftMegarepoTarget, MegarepoError> {
        let repo = mononoke
            .repo_name_from_id(RepositoryId::new(self.repo_id as i32))
            .map(|name| ThriftRepoSpecifier {
                name,
                ..Default::default()
            });
        Ok(ThriftMegarepoTarget {
            repo_id: Some(self.repo_id),
            bookmark: self.bookmark,
            repo,
            ..Default::default()
        })
    }
}

#[async_trait]
impl IntoApiFormat<ThriftMegarepoSyncTargetConfig> for SyncTargetConfig {
    fn into_api_format(
        self,
        mononoke: &Mononoke,
    ) -> Result<ThriftMegarepoSyncTargetConfig, MegarepoError> {
        Ok(ThriftMegarepoSyncTargetConfig {
            target: self.target.into_api_format(mononoke)?,
            sources: self.sources,
            version: self.version,
            ..Default::default()
        })
    }
}

#[cfg(test)]
mod test {
    use blobstore::Loadable;
    use blobstore::PutBehaviour;
    use blobstore::Storable;
    use context::CoreContext;
    use fbinit::FacebookInit;
    use memblob::Memblob;

    use super::*;

    macro_rules! test_blobstore_key {
        {
            $type: ident,
            $prefix: expr
        } => {
            let id = $type::from_byte_array([1; 32]);
            assert_eq!(id.blobstore_key(), format!(concat!($prefix, ".blake2.{}"), id));
        }
    }

    macro_rules! serialize_deserialize {
        {
            $type: ident
        } => {
            let id = $type::from_byte_array([1; 32]);
            let serialized = serde_json::to_string(&id).unwrap();
            let deserialized = serde_json::from_str(&serialized).unwrap();
            assert_eq!(id, deserialized);
        }
    }

    #[test]
    fn blobstore_key() {
        // These IDs are persistent, and this test is really to make sure that they don't change
        // accidentally. Same as in typed_hash.rs
        test_blobstore_key!(
            MegarepoAsynchronousRequestParamsId,
            "async.svc.MegarepoAsynchronousRequestParams"
        );
        test_blobstore_key!(
            MegarepoAsynchronousRequestResultId,
            "async.svc.MegarepoAsynchronousRequestResult"
        );
        test_blobstore_key!(
            MegarepoAsynchronousRequestParamsId,
            "async.svc.MegarepoAsynchronousRequestParams"
        );
        test_blobstore_key!(
            MegarepoAsynchronousRequestResultId,
            "async.svc.MegarepoAsynchronousRequestResult"
        );
    }

    #[test]
    fn test_serialize_deserialize() {
        serialize_deserialize!(MegarepoAsynchronousRequestParamsId);
        serialize_deserialize!(MegarepoAsynchronousRequestResultId);
    }

    macro_rules! test_store_load {
        { $type: ident, $ctx: ident, $blobstore: ident } => {
            let obj = $type::from_thrift(Default::default());

            let id = obj
                .clone()
                .store(&$ctx, &$blobstore)
                .await
                .expect(&format!("Failed to store {}", stringify!($type)));

            let obj2 = id
                .load(&$ctx, &$blobstore)
                .await
                .expect(&format!("Failed to load {}", stringify!($type)));

            assert_eq!(obj, obj2);
        }
    }

    #[fbinit::test]
    async fn test_megaerpo_add_target_params_type(fb: FacebookInit) {
        let blobstore = Memblob::new(PutBehaviour::IfAbsent);
        let ctx = CoreContext::test_mock(fb);
        test_store_load!(MegarepoAsynchronousRequestParams, ctx, blobstore);
        test_store_load!(MegarepoAsynchronousRequestResult, ctx, blobstore);
    }
}


/*
 * 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 methods;
mod worker;

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

use anyhow::Error;
use clap::Parser;
use cmdlib_logging::ScribeLoggingArgs;
use context::SessionContainer;
use environment::WarmBookmarksCacheDerivedData;
use fbinit::FacebookInit;
use hostname::get_hostname;
use megarepo_api::MegarepoApi;
use metaconfig_types::ShardedService;
use mononoke_app::args::HooksAppExtension;
use mononoke_app::args::RepoFilterAppExtension;
use mononoke_app::args::ShutdownTimeoutArgs;
use mononoke_app::args::WarmBookmarksCacheExtension;
use mononoke_app::fb303::AliveService;
use mononoke_app::fb303::Fb303AppExtension;
use mononoke_app::MononokeAppBuilder;

const SERVICE_NAME: &str = "megarepo_async_requests_worker";

/// Processes the megarepo async requests
#[derive(Parser)]
struct AsyncRequestsWorkerArgs {
    #[clap(flatten)]
    shutdown_timeout_args: ShutdownTimeoutArgs,
    #[clap(flatten)]
    scribe_logging_args: ScribeLoggingArgs,
    /// The number of requests to process before exiting
    #[clap(long)]
    request_limit: Option<usize>,
    /// The number of requests / jobs to be processed concurrently
    #[clap(long, short = 'j', default_value = "1")]
    jobs: usize,
}

#[fbinit::main]
fn main(fb: FacebookInit) -> Result<(), Error> {
    let app = MononokeAppBuilder::new(fb)
        .with_warm_bookmarks_cache(WarmBookmarksCacheDerivedData::NoDerivation)
        .with_app_extension(WarmBookmarksCacheExtension {})
        .with_app_extension(HooksAppExtension {})
        .with_app_extension(Fb303AppExtension {})
        .with_app_extension(RepoFilterAppExtension {})
        .build::<AsyncRequestsWorkerArgs>()?;
    let args: AsyncRequestsWorkerArgs = app.args()?;
    let request_limit = args.request_limit;
    let jobs_limit = args.jobs;
    let (env, logger, runtime) = (app.environment(), app.logger(), app.runtime());

    let session = SessionContainer::new_with_defaults(env.fb);
    let ctx = session.new_context(logger.clone(), env.scuba_sample_builder.clone());

    let mononoke = Arc::new(
        runtime
            .block_on(app.open_managed_repos(Some(ShardedService::AsyncRequestsWorker)))?
            .make_mononoke_api()?,
    );
    let megarepo = Arc::new(runtime.block_on(MegarepoApi::new(&app, mononoke))?);

    let tw_job_cluster = std::env::var("TW_JOB_CLUSTER");
    let tw_job_name = std::env::var("TW_JOB_NAME");
    let tw_task_id = std::env::var("TW_TASK_ID");

    let name = match (tw_job_cluster, tw_job_name, tw_task_id) {
        (Ok(tw_job_cluster), Ok(tw_job_name), Ok(tw_task_id)) => {
            format!("{}/{}/{}", tw_job_cluster, tw_job_name, tw_task_id)
        }
        _ => format!(
            "megarepo_async_requests_worker/{}",
            get_hostname().unwrap_or_else(|_| "unknown_hostname".to_string())
        ),
    };

    let will_exit = Arc::new(AtomicBool::new(false));
    let worker = worker::AsyncMethodRequestWorker::new(megarepo, name);

    app.start_monitoring(SERVICE_NAME, AliveService)?;
    app.start_stats_aggregation()?;

    let run_worker = {
        let will_exit = will_exit.clone();
        move |_app| async move {
            Ok(worker
                .run(&ctx, will_exit, request_limit, jobs_limit)
                .await?)
        }
    };

    app.run_until_terminated(
        run_worker,
        move || will_exit.store(true, Ordering::Relaxed),
        args.shutdown_timeout_args.shutdown_grace_period,
        async {
            // the code to gracefully stop things goes here
        },
        args.shutdown_timeout_args.shutdown_timeout,
    )?;

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

//! The concrete method implementations.
//!
//! This module provides the concrete implementations for methods - functions
//! taking the Params and returning the Results - to be used by the request worker.
//! This module is not aware of the async nature of those methods. All the token
//! handling, enqueuing and polling should be done by the callers.

use std::collections::HashMap;

use anyhow::anyhow;
use async_requests::types::IntoConfigFormat;
use async_requests::types::MegarepoAsynchronousRequestParams;
use async_requests::types::MegarepoAsynchronousRequestResult;
use context::CoreContext;
use megarepo_api::MegarepoApi;
use megarepo_error::MegarepoError;
use mononoke_types::ChangesetId;
use source_control as thrift;

async fn megarepo_sync_changeset(
    ctx: &CoreContext,
    megarepo_api: &MegarepoApi,
    params: thrift::MegarepoSyncChangesetParams,
) -> Result<thrift::MegarepoSyncChangesetResponse, MegarepoError> {
    let source_cs_id = ChangesetId::from_bytes(params.cs_id).map_err(MegarepoError::request)?;
    let target_location =
        ChangesetId::from_bytes(params.target_location).map_err(MegarepoError::request)?;
    let cs_id = megarepo_api
        .sync_changeset(
            ctx,
            source_cs_id,
            params.source_name,
            params.target.into_config_format(&megarepo_api.mononoke())?,
            target_location,
        )
        .await?
        .as_ref()
        .into();
    Ok(thrift::MegarepoSyncChangesetResponse {
        cs_id,
        ..Default::default()
    })
}

async fn megarepo_add_sync_target(
    ctx: &CoreContext,
    megarepo_api: &MegarepoApi,
    params: thrift::MegarepoAddTargetParams,
) -> Result<thrift::MegarepoAddTargetResponse, MegarepoError> {
    let config = params.config_with_new_target;
    let mut changesets_to_merge = HashMap::new();
    for (s, cs_id) in params.changesets_to_merge {
        let cs_id = ChangesetId::from_bytes(cs_id).map_err(MegarepoError::request)?;
        changesets_to_merge.insert(s, cs_id);
    }
    let cs_id = megarepo_api
        .add_sync_target(
            ctx,
            config.into_config_format(&megarepo_api.mononoke())?,
            changesets_to_merge,
            params.message,
        )
        .await?
        .as_ref()
        .into();
    Ok(thrift::MegarepoAddTargetResponse {
        cs_id,
        ..Default::default()
    })
}

async fn megarepo_add_branching_sync_target(
    ctx: &CoreContext,
    megarepo_api: &MegarepoApi,
    params: thrift::MegarepoAddBranchingTargetParams,
) -> Result<thrift::MegarepoAddBranchingTargetResponse, MegarepoError> {
    let cs_id = megarepo_api
        .add_branching_sync_target(
            ctx,
            params.target.into_config_format(&megarepo_api.mononoke())?,
            ChangesetId::from_bytes(params.branching_point).map_err(MegarepoError::request)?,
            params
                .source_target
                .into_config_format(&megarepo_api.mononoke())?,
        )
        .await?
        .as_ref()
        .into();
    Ok(thrift::MegarepoAddBranchingTargetResponse {
        cs_id,
        ..Default::default()
    })
}

async fn megarepo_change_target_config(
    ctx: &CoreContext,
    megarepo_api: &MegarepoApi,
    params: thrift::MegarepoChangeTargetConfigParams,
) -> Result<thrift::MegarepoChangeTargetConfigResponse, MegarepoError> {
    let mut changesets_to_merge = HashMap::new();
    for (s, cs_id) in params.changesets_to_merge {
        let cs_id = ChangesetId::from_bytes(cs_id).map_err(MegarepoError::request)?;
        changesets_to_merge.insert(s, cs_id);
    }
    let target_location =
        ChangesetId::from_bytes(params.target_location).map_err(MegarepoError::request)?;
    let cs_id = megarepo_api
        .change_target_config(
            ctx,
            params.target.into_config_format(&megarepo_api.mononoke())?,
            params.new_version,
            target_location,
            changesets_to_merge,
            params.message,
        )
        .await?
        .as_ref()
        .into();
    Ok(thrift::MegarepoChangeTargetConfigResponse {
        cs_id,
        ..Default::default()
    })
}

async fn megarepo_remerge_source(
    ctx: &CoreContext,
    megarepo_api: &MegarepoApi,
    params: thrift::MegarepoRemergeSourceParams,
) -> Result<thrift::MegarepoRemergeSourceResponse, MegarepoError> {
    let remerge_cs_id = ChangesetId::from_bytes(params.cs_id).map_err(MegarepoError::request)?;
    let target_location =
        ChangesetId::from_bytes(params.target_location).map_err(MegarepoError::request)?;
    let cs_id = megarepo_api
        .remerge_source(
            ctx,
            params.source_name,
            remerge_cs_id,
            params.message,
            &params.target.into_config_format(&megarepo_api.mononoke())?,
            target_location,
        )
        .await?
        .as_ref()
        .into();

    Ok(thrift::MegarepoRemergeSourceResponse {
        cs_id,
        ..Default::default()
    })
}

/// Given the request params dispatches the request to the right processing
/// funtion and returns the computation result. This function doesn't return
/// `Result` as both successfull computation and error are part of
/// `MegarepoAsynchronousRequestResult` structure.
pub(crate) async fn megarepo_async_request_compute(
    ctx: &CoreContext,
    megarepo_api: &MegarepoApi,
    params: MegarepoAsynchronousRequestParams,
) -> MegarepoAsynchronousRequestResult {
    match params.into() {
        megarepo_types_thrift::MegarepoAsynchronousRequestParams::megarepo_add_target_params(params) => {
            megarepo_add_sync_target(ctx, megarepo_api, params)
                .await
                .into()
        }
        megarepo_types_thrift::MegarepoAsynchronousRequestParams::megarepo_add_branching_target_params(params) => {
            megarepo_add_branching_sync_target(ctx, megarepo_api, params)
                .await
                .into()
        }
        megarepo_types_thrift::MegarepoAsynchronousRequestParams::megarepo_change_target_params(params) => {
            megarepo_change_target_config(ctx, megarepo_api, params)
                .await
                .into()
        }
        megarepo_types_thrift::MegarepoAsynchronousRequestParams::megarepo_remerge_source_params(params) => {
            megarepo_remerge_source(ctx, megarepo_api, params)
                .await
                .into()
        }
        megarepo_types_thrift::MegarepoAsynchronousRequestParams::megarepo_sync_changeset_params(params) => {
            megarepo_sync_changeset(ctx, megarepo_api, params)
                .await
                .into()
        }
        megarepo_types_thrift::MegarepoAsynchronousRequestParams::UnknownField(union_tag) => {
            Err::<thrift::MegarepoRemergeSourceResponse, _>(MegarepoError::internal(anyhow!(
                "this type of reuqest (MegarepoAsynchronousRequestParams tag {}) not supported by this worker!", union_tag
            )))
            .into()

        }
    }
}


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

//! AsyncMethodRequestWorker is an object that provides everything that's needed
//! for processing the requests from the queue.
//!
//! In can grab requests from the queue, compute the result and update the
//! requests table with a response.
//! One important consideration to keep in mind - worker executes request "at least once"
//! but not exactly once i.e. the same request might be executed a few times.

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

use async_requests::types::IntoConfigFormat;
use async_requests::types::MegarepoAsynchronousRequestParams;
use async_requests::AsyncMethodRequestQueue;
use async_requests::ClaimedBy;
use async_requests::RequestId;
use async_stream::try_stream;
use cloned::cloned;
use context::CoreContext;
use futures::future::abortable;
use futures::future::select;
use futures::future::Either;
use futures::pin_mut;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures::Stream;
use megarepo_api::MegarepoApi;
use megarepo_config::Target;
use megarepo_error::MegarepoError;
use mononoke_types::RepositoryId;
use mononoke_types::Timestamp;
use slog::debug;
use slog::error;
use slog::info;

use crate::methods::megarepo_async_request_compute;

const DEQUEUE_STREAM_SLEEP_TIME: u64 = 1000;
// Number of seconds after which inprogress request is considered abandoned
// if it hasn't updated inprogress timestamp
const ABANDONED_REQUEST_THRESHOLD_SECS: i64 = 5 * 60;
const KEEP_ALIVE_INTERVAL: Duration = Duration::from_secs(10);

#[derive(Clone)]
pub struct AsyncMethodRequestWorker {
    megarepo: Arc<MegarepoApi>,
    name: String,
}

impl AsyncMethodRequestWorker {
    /// Creates a new tailer instance that's going to use provided megarepo API
    /// The name argument should uniquely identify tailer instance and will be put
    /// in the queue table so it's possible to find out which instance is working on
    /// a given task (for debugging purposes).
    pub fn new(megarepo: Arc<MegarepoApi>, name: String) -> Self {
        Self { megarepo, name }
    }

    /// Start async request worker.
    /// If limit is set the worker will process a preset number of requests and
    /// return. If the limit is None the worker will be running continuously. The
    /// will_exit atomic bool is a flag to prevent the worker from grabbing new
    /// items from the queue and gracefully terminate.
    pub async fn run(
        &self,
        ctx: &CoreContext,
        will_exit: Arc<AtomicBool>,
        limit: Option<usize>,
        concurrency_limit: usize,
    ) -> Result<(), MegarepoError> {
        let queues_with_repos = self.megarepo.all_async_method_request_queues(ctx).await?;

        // Build stream that pools all the queues
        let request_stream = self
            .request_stream(ctx.clone(), queues_with_repos, will_exit)
            .boxed();

        let request_stream = if let Some(limit) = limit {
            request_stream.take(limit).left_stream()
        } else {
            request_stream.right_stream()
        };

        info!(
            ctx.logger(),
            "Worker initialization complete, starting request processing loop.",
        );

        request_stream
            .try_for_each_concurrent(Some(concurrency_limit), async move |(req_id, params)| {
                let worker = self.clone();
                let ctx = ctx.clone();
                let _updated = tokio::spawn(worker.compute_and_mark_completed(ctx, req_id, params))
                    .await
                    .map_err(MegarepoError::internal)??;
                Ok(())
            })
            .await?;
        Ok(())
    }

    pub fn request_stream(
        &self,
        ctx: CoreContext,
        queues_with_repos: Vec<(Vec<RepositoryId>, AsyncMethodRequestQueue)>,
        will_exit: Arc<AtomicBool>,
    ) -> impl Stream<Item = Result<(RequestId, MegarepoAsynchronousRequestParams), MegarepoError>>
    {
        let claimed_by = ClaimedBy(self.name.clone());
        let sleep_time = Duration::from_millis(DEQUEUE_STREAM_SLEEP_TIME);
        Self::request_stream_inner(
            ctx,
            claimed_by,
            queues_with_repos,
            will_exit,
            sleep_time,
            ABANDONED_REQUEST_THRESHOLD_SECS,
        )
    }

    fn request_stream_inner(
        ctx: CoreContext,
        claimed_by: ClaimedBy,
        queues_with_repos: Vec<(Vec<RepositoryId>, AsyncMethodRequestQueue)>,
        will_exit: Arc<AtomicBool>,
        sleep_time: Duration,
        abandoned_threshold_secs: i64,
    ) -> impl Stream<Item = Result<(RequestId, MegarepoAsynchronousRequestParams), MegarepoError>>
    {
        try_stream! {
            'outer: loop {
                let mut yielded = false;
                for (repo_ids, queue) in &queues_with_repos {
                    Self::cleanup_abandoned_requests(
                        &ctx,
                        repo_ids,
                        queue,
                        abandoned_threshold_secs
                    ).await?;
                    if will_exit.load(Ordering::Relaxed) {
                        break 'outer;
                    }
                    if let Some((request_id, params)) = queue.dequeue(&ctx, &claimed_by, repo_ids).await? {
                        yield (request_id, params);
                        yielded = true;
                    }
                }
                if ! yielded {
                    // No requests in the queues, sleep before trying again.
                    debug!(
                        ctx.logger(),
                        "nothing to do, sleeping",
                    );
                    tokio::time::sleep(sleep_time).await;

                }
            }
        }
    }

    async fn cleanup_abandoned_requests(
        ctx: &CoreContext,
        repo_ids: &[RepositoryId],
        queue: &AsyncMethodRequestQueue,
        abandoned_threshold_secs: i64,
    ) -> Result<(), MegarepoError> {
        let now = Timestamp::now();
        let abandoned_timestamp =
            Timestamp::from_timestamp_secs(now.timestamp_seconds() - abandoned_threshold_secs);
        let requests = queue
            .find_abandoned_requests(ctx, repo_ids, abandoned_timestamp)
            .await?;
        if !requests.is_empty() {
            ctx.scuba().clone().log_with_msg(
                "Find requests to abandon",
                Some(format!("{}", requests.len())),
            );
        }

        for req_id in requests {
            if queue
                .mark_abandoned_request_as_new(ctx, req_id.clone(), abandoned_timestamp)
                .await?
            {
                ctx.scuba()
                    .clone()
                    .add("request_id", req_id.0.0)
                    .log_with_msg("Abandoned request", None);
            }
        }
        Ok(())
    }

    /// Params into stored response. Doesn't mark it as "in progress" (as this is done during dequeueing).
    /// Returns true if the result was successfully stored. Returns false if we
    /// lost the race (the request table was updated).
    async fn compute_and_mark_completed(
        self,
        ctx: CoreContext,
        req_id: RequestId,
        params: MegarepoAsynchronousRequestParams,
    ) -> Result<bool, MegarepoError> {
        let target = params
            .target()?
            .clone()
            .into_config_format(&self.megarepo.mononoke())?;
        let queue = self
            .megarepo
            .async_method_request_queue(&ctx, &target)
            .await?;

        let ctx = self.prepare_ctx(&ctx, &req_id, &target);

        // Do the actual work.
        let work_fut = megarepo_async_request_compute(&ctx, &self.megarepo, params);

        // Start the loop that would keep saying that request is still being
        // processed
        let (keep_alive, keep_alive_abort_handle) = abortable({
            cloned!(ctx, req_id, queue);
            async move { Self::keep_alive_loop(&ctx, &req_id, &queue).await }
        });

        let keep_alive = tokio::spawn(keep_alive);

        pin_mut!(work_fut);
        pin_mut!(keep_alive);
        match select(work_fut, keep_alive).await {
            Either::Left((result, _)) => {
                // We completed the request - let's mark it as complete
                keep_alive_abort_handle.abort();
                info!(
                    ctx.logger(),
                    "[{}] request complete, saving result", &req_id.0
                );
                ctx.scuba()
                    .clone()
                    .log_with_msg("Request complete, saving result", None);

                // Save the result.
                let updated_res = queue.complete(&ctx, &req_id, result).await;

                let updated = match updated_res {
                    Ok(updated) => {
                        info!(ctx.logger(), "[{}] result saved", &req_id.0);
                        ctx.scuba().clone().log_with_msg("Result saved", None);
                        updated
                    }
                    Err(err) => {
                        ctx.scuba()
                            .clone()
                            .log_with_msg("Failed to save result", Some(format!("{:?}", err)));
                        return Err(err);
                    }
                };

                Ok(updated)
            }
            Either::Right((res, _)) => {
                // We haven't completed the request, and failed to update
                // inprogress timestamp. Most likely it means that other
                // worker has completed it

                res.map_err(MegarepoError::internal)?
                    .map_err(MegarepoError::internal)?;
                info!(
                    ctx.logger(),
                    "[{}] was completed by other worker, stopping", &req_id.0
                );
                Ok(false)
            }
        }
    }

    fn prepare_ctx(&self, ctx: &CoreContext, req_id: &RequestId, target: &Target) -> CoreContext {
        let ctx = ctx.with_mutated_scuba(|mut scuba| {
            scuba.add("request_id", req_id.0.0);
            scuba
        });

        info!(
            ctx.logger(),
            "[{}] new request:  id: {}, type: {}, repo_id: {}, bookmark: {}",
            &req_id.0,
            &req_id.0,
            &req_id.1,
            &target.repo_id,
            &target.bookmark,
        );

        ctx
    }

    async fn keep_alive_loop(
        ctx: &CoreContext,
        req_id: &RequestId,
        queue: &AsyncMethodRequestQueue,
    ) {
        loop {
            let mut scuba = ctx.scuba().clone();
            ctx.perf_counters().insert_perf_counters(&mut scuba);

            let res = queue.update_in_progress_timestamp(ctx, req_id).await;
            match res {
                Ok(res) => {
                    // Weren't able to update inprogress timestamp - that probably means
                    // that request was completed by someone else. Exiting
                    if !res {
                        scuba.log_with_msg(
                            "Race while updating inprogress timestamp, exiting keep-alive loop",
                            None,
                        );
                        break;
                    }
                    scuba.log_with_msg("Updated inprogress timestamp", None);
                }
                Err(err) => {
                    error!(
                        ctx.logger(),
                        "[{}] failed to update inprogress timestamp: {:?}", req_id.0, err
                    );
                    scuba.log_with_msg(
                        "Failed to update inprogress timestamp",
                        Some(format!("{:?}", err)),
                    );
                }
            }
            tokio::time::sleep(KEEP_ALIVE_INTERVAL).await;
        }
    }
}

#[cfg(test)]
mod test {
    use std::sync::atomic::Ordering;

    use anyhow::Error;
    use fbinit::FacebookInit;
    use mononoke_api::Mononoke;
    use mononoke_api::Repo;
    use requests_table::RequestType;
    use source_control as thrift;

    use super::*;

    #[fbinit::test]
    async fn test_request_stream_simple(fb: FacebookInit) -> Result<(), Error> {
        let q = AsyncMethodRequestQueue::new_test_in_memory().unwrap();
        let ctx = CoreContext::test_mock(fb);
        let repo: Repo = test_repo_factory::build_empty(ctx.fb).await?;
        let mononoke = Mononoke::new_test(vec![("test".to_string(), repo)]).await?;

        let params = thrift::MegarepoSyncChangesetParams {
            cs_id: vec![],
            source_name: "name".to_string(),
            target: thrift::MegarepoTarget {
                repo_id: Some(0),
                bookmark: "book".to_string(),
                ..Default::default()
            },
            target_location: vec![],
            ..Default::default()
        };
        q.enqueue(ctx.clone(), &mononoke, params).await?;

        let will_exit = Arc::new(AtomicBool::new(false));
        let s = AsyncMethodRequestWorker::request_stream_inner(
            ctx,
            ClaimedBy("name".to_string()),
            vec![(vec![RepositoryId::new(0)], q)],
            will_exit.clone(),
            Duration::from_millis(100),
            ABANDONED_REQUEST_THRESHOLD_SECS,
        );

        let s = tokio::spawn(s.try_collect::<Vec<_>>());
        tokio::time::sleep(Duration::from_secs(1)).await;
        will_exit.store(true, Ordering::Relaxed);
        let res = s.await??;
        assert_eq!(res.len(), 1);
        assert_eq!(
            res[0].0.1,
            RequestType("megarepo_sync_changeset".to_string())
        );
        Ok(())
    }

    #[fbinit::test]
    async fn test_request_stream_clear_abandoned(fb: FacebookInit) -> Result<(), Error> {
        let q = AsyncMethodRequestQueue::new_test_in_memory().unwrap();
        let ctx = CoreContext::test_mock(fb);
        let repo: Repo = test_repo_factory::build_empty(ctx.fb).await?;
        let mononoke = Mononoke::new_test(vec![("test".to_string(), repo)]).await?;

        let params = thrift::MegarepoSyncChangesetParams {
            cs_id: vec![],
            source_name: "name".to_string(),
            target: thrift::MegarepoTarget {
                repo_id: Some(0),
                bookmark: "book".to_string(),
                ..Default::default()
            },
            target_location: vec![],
            ..Default::default()
        };
        q.enqueue(ctx.clone(), &mononoke, params).await?;

        // Grab it from the queue...
        let dequed = q
            .dequeue(
                &ctx,
                &ClaimedBy("name".to_string()),
                &[RepositoryId::new(0)],
            )
            .await?;
        assert!(dequed.is_some());

        // ... and check that the queue is empty now...
        let will_exit = Arc::new(AtomicBool::new(false));
        let s = AsyncMethodRequestWorker::request_stream_inner(
            ctx.clone(),
            ClaimedBy("name".to_string()),
            vec![(vec![RepositoryId::new(0)], q.clone())],
            will_exit.clone(),
            Duration::from_millis(100),
            ABANDONED_REQUEST_THRESHOLD_SECS,
        );

        let s = tokio::spawn(s.try_collect::<Vec<_>>());
        tokio::time::sleep(Duration::from_secs(1)).await;
        will_exit.store(true, Ordering::Relaxed);
        let res = s.await??;
        assert_eq!(res, vec![]);

        // ... now make it "abandoned", and make sure we reclaim it
        tokio::time::sleep(Duration::from_secs(1)).await;
        let will_exit = Arc::new(AtomicBool::new(false));
        let s = AsyncMethodRequestWorker::request_stream_inner(
            ctx,
            ClaimedBy("name".to_string()),
            vec![(vec![RepositoryId::new(0)], q)],
            will_exit.clone(),
            Duration::from_millis(100),
            1, // 1 second
        );

        let s = tokio::spawn(s.try_collect::<Vec<_>>());
        tokio::time::sleep(Duration::from_secs(1)).await;
        will_exit.store(true, Ordering::Relaxed);
        let res = s.await??;
        assert_eq!(res.len(), 1);

        Ok(())
    }
}