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

use std::collections::HashMap;
use std::num::NonZeroU64;
use std::time::Duration;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use async_trait::async_trait;
use blobstore::Loadable;
use edenapi_types::wire::WireCommitHashToLocationRequestBatch;
use edenapi_types::AlterSnapshotRequest;
use edenapi_types::AlterSnapshotResponse;
use edenapi_types::AnyFileContentId;
use edenapi_types::AnyId;
use edenapi_types::Batch;
use edenapi_types::BonsaiFileChange;
use edenapi_types::CommitGraphEntry;
use edenapi_types::CommitGraphRequest;
use edenapi_types::CommitGraphSegmentParent;
use edenapi_types::CommitGraphSegmentsEntry;
use edenapi_types::CommitGraphSegmentsRequest;
use edenapi_types::CommitHashLookupRequest;
use edenapi_types::CommitHashLookupResponse;
use edenapi_types::CommitHashToLocationResponse;
use edenapi_types::CommitId;
use edenapi_types::CommitIdScheme;
use edenapi_types::CommitLocationToHashRequest;
use edenapi_types::CommitLocationToHashRequestBatch;
use edenapi_types::CommitLocationToHashResponse;
use edenapi_types::CommitMutationsRequest;
use edenapi_types::CommitMutationsResponse;
use edenapi_types::CommitRevlogData;
use edenapi_types::CommitRevlogDataRequest;
use edenapi_types::CommitTranslateIdRequest;
use edenapi_types::CommitTranslateIdResponse;
use edenapi_types::EphemeralPrepareRequest;
use edenapi_types::EphemeralPrepareResponse;
use edenapi_types::FetchSnapshotRequest;
use edenapi_types::FetchSnapshotResponse;
use edenapi_types::UploadBonsaiChangesetRequest;
use edenapi_types::UploadHgChangesetsRequest;
use edenapi_types::UploadToken;
use edenapi_types::UploadTokensResponse;
use ephemeral_blobstore::BubbleId;
use futures::stream;
use futures::try_join;
use futures::Stream;
use futures::StreamExt;
use futures::TryStreamExt;
use gotham::state::FromState;
use gotham::state::State;
use gotham_derive::StateData;
use gotham_derive::StaticResponseExtender;
use gotham_ext::error::HttpError;
use gotham_ext::middleware::scuba::ScubaMiddlewareState;
use gotham_ext::response::TryIntoResponse;
use mercurial_types::HgChangesetId;
use mercurial_types::HgNodeHash;
use mononoke_api::CreateInfo;
use mononoke_api::MononokeError;
use mononoke_api_hg::HgRepoContext;
use mononoke_types::hash::GitSha1;
use mononoke_types::ChangesetId;
use mononoke_types::DateTime;
use mononoke_types::FileChange;
use mononoke_types::Globalrev;
use serde::Deserialize;
use tunables::tunables;
use types::HgId;
use types::Parents;

use super::handler::EdenApiContext;
use super::EdenApiHandler;
use super::EdenApiMethod;
use super::HandlerInfo;
use super::HandlerResult;
use crate::context::ServerContext;
use crate::errors::ErrorKind;
use crate::middleware::request_dumper::RequestDumper;
use crate::middleware::RequestContext;
use crate::utils::cbor_stream_filtered_errors;
use crate::utils::custom_cbor_stream;
use crate::utils::get_repo;
use crate::utils::parse_cbor_request;
use crate::utils::parse_wire_request;
use crate::utils::to_create_change;
use crate::utils::to_hg_path;
use crate::utils::to_mononoke_path;
use crate::utils::to_revlog_changeset;

/// XXX: This number was chosen arbitrarily.
const MAX_CONCURRENT_FETCHES_PER_REQUEST: usize = 100;
const HASH_TO_LOCATION_BATCH_SIZE: usize = 100;

const PHASES_CHECK_LIMIT: usize = 10;

#[derive(Debug, Deserialize, StateData, StaticResponseExtender)]
pub struct HashToLocationParams {
    repo: String,
}

#[derive(Debug, Deserialize, StateData, StaticResponseExtender)]
pub struct RevlogDataParams {
    repo: String,
}

#[derive(Debug, Deserialize, StateData, StaticResponseExtender)]
pub struct UploadBonsaiChangesetQueryString {
    bubble_id: Option<NonZeroU64>,
}

pub struct LocationToHashHandler;

async fn translate_location(
    hg_repo_ctx: HgRepoContext,
    request: CommitLocationToHashRequest,
) -> Result<CommitLocationToHashResponse, Error> {
    let location = request.location.map_descendant(|x| x.into());
    let ancestors: Vec<HgChangesetId> = hg_repo_ctx
        .location_to_hg_changeset_id(location, request.count)
        .await
        .context(ErrorKind::CommitLocationToHashRequestFailed)?;
    let hgids = ancestors.into_iter().map(|x| x.into()).collect();
    let answer = CommitLocationToHashResponse {
        location: request.location,
        count: request.count,
        hgids,
    };
    Ok(answer)
}

#[async_trait]
impl EdenApiHandler for LocationToHashHandler {
    type Request = CommitLocationToHashRequestBatch;
    type Response = CommitLocationToHashResponse;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::CommitLocationToHash;
    const ENDPOINT: &'static str = "/commit/location_to_hash";

    fn sampling_rate(_request: &Self::Request) -> NonZeroU64 {
        nonzero_ext::nonzero!(100u64)
    }

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        let hgid_list = request
            .requests
            .into_iter()
            .map(move |location| translate_location(repo.clone(), location));
        let response = stream::iter(hgid_list).buffer_unordered(MAX_CONCURRENT_FETCHES_PER_REQUEST);
        Ok(response.boxed())
    }
}

pub async fn hash_to_location(state: &mut State) -> Result<impl TryIntoResponse, HttpError> {
    async fn hash_to_location_chunk(
        hg_repo_ctx: HgRepoContext,
        master_heads: Vec<HgChangesetId>,
        hg_cs_ids: Vec<HgChangesetId>,
    ) -> impl Stream<Item = CommitHashToLocationResponse> {
        let hgcsid_to_location = hg_repo_ctx
            .many_changeset_ids_to_locations(master_heads, hg_cs_ids.clone())
            .await;
        let responses = hg_cs_ids.into_iter().map(move |hgcsid| {
            let result = match hgcsid_to_location.as_ref() {
                Ok(hsh) => match hsh.get(&hgcsid) {
                    Some(Ok(l)) => Ok(Some(l.map_descendant(|x| x.into()))),
                    Some(Err(e)) => Err(e.into()),
                    None => Ok(None),
                },
                Err(e) => Err(e.into()),
            };
            CommitHashToLocationResponse {
                hgid: hgcsid.into(),
                result,
            }
        });
        stream::iter(responses)
    }

    let params = HashToLocationParams::take_from(state);

    state.put(HandlerInfo::new(
        &params.repo,
        EdenApiMethod::CommitHashToLocation,
    ));

    ScubaMiddlewareState::try_set_sampling_rate(state, nonzero_ext::nonzero!(256_u64));

    let sctx = ServerContext::borrow_from(state);
    let rctx = RequestContext::borrow_from(state).clone();

    let hg_repo_ctx = get_repo(sctx, &rctx, &params.repo, None).await?;

    let batch = parse_wire_request::<WireCommitHashToLocationRequestBatch>(state).await?;

    if let Some(rd) = RequestDumper::try_borrow_mut_from(state) {
        rd.add_request(&batch);
    };

    let master_heads = batch
        .master_heads
        .into_iter()
        .map(|x| x.into())
        .collect::<Vec<_>>();

    let response = stream::iter(batch.hgids)
        .chunks(HASH_TO_LOCATION_BATCH_SIZE)
        .map(|chunk| chunk.into_iter().map(|x| x.into()).collect::<Vec<_>>())
        .map({
            let ctx = hg_repo_ctx.clone();
            move |chunk| hash_to_location_chunk(ctx.clone(), master_heads.clone(), chunk)
        })
        .buffer_unordered(3)
        .flatten();
    let cbor_response = custom_cbor_stream(super::monitor_request(state, response), |t| {
        t.result.as_ref().err()
    });
    Ok(cbor_response)
}

pub async fn revlog_data(state: &mut State) -> Result<impl TryIntoResponse, HttpError> {
    let params = RevlogDataParams::take_from(state);

    state.put(HandlerInfo::new(
        &params.repo,
        EdenApiMethod::CommitRevlogData,
    ));

    let sctx = ServerContext::borrow_from(state);
    let rctx = RequestContext::borrow_from(state).clone();

    let hg_repo_ctx = get_repo(sctx, &rctx, &params.repo, None).await?;

    let request: CommitRevlogDataRequest = parse_cbor_request(state).await?;
    let revlog_commits = request
        .hgids
        .into_iter()
        .map(move |hg_id| commit_revlog_data(hg_repo_ctx.clone(), hg_id));
    let response =
        stream::iter(revlog_commits).buffer_unordered(MAX_CONCURRENT_FETCHES_PER_REQUEST);
    Ok(cbor_stream_filtered_errors(super::monitor_request(
        state, response,
    )))
}

async fn commit_revlog_data(
    hg_repo_ctx: HgRepoContext,
    hg_id: HgId,
) -> Result<CommitRevlogData, Error> {
    let bytes = hg_repo_ctx
        .revlog_commit_data(hg_id.into())
        .await
        .context(ErrorKind::CommitRevlogDataRequestFailed)?
        .ok_or(ErrorKind::HgIdNotFound(hg_id))?;
    let answer = CommitRevlogData::new(hg_id, bytes);
    Ok(answer)
}

pub struct HashLookupHandler;

#[async_trait]
impl EdenApiHandler for HashLookupHandler {
    type Request = Batch<CommitHashLookupRequest>;
    type Response = CommitHashLookupResponse;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::CommitHashLookup;
    const ENDPOINT: &'static str = "/commit/hash_lookup";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        use CommitHashLookupRequest::*;
        Ok(stream::iter(request.batch.into_iter())
            .then(move |request| {
                let hg_repo_ctx = repo.clone();
                async move {
                    let changesets = match request {
                        InclusiveRange(low, high) => {
                            hg_repo_ctx.get_hg_in_range(low.into(), high.into()).await?
                        }
                    };
                    let hgids = changesets.into_iter().map(|x| x.into()).collect();
                    let response = CommitHashLookupResponse { request, hgids };
                    Ok(response)
                }
            })
            .boxed())
    }
}

/// Upload list of HgChangesets requested by the client
pub struct UploadHgChangesetsHandler;

#[async_trait]
impl EdenApiHandler for UploadHgChangesetsHandler {
    type Request = UploadHgChangesetsRequest;
    type Response = UploadTokensResponse;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::UploadHgChangesets;
    const ENDPOINT: &'static str = "/upload/changesets";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        let changesets = request.changesets;
        let mutations = request.mutations;
        let changesets_data = changesets
            .into_iter()
            .map(|changeset| {
                Ok((
                    HgChangesetId::new(HgNodeHash::from(changeset.node_id)),
                    to_revlog_changeset(changeset.changeset_content)?,
                ))
            })
            .collect::<Result<Vec<_>, Error>>()?;

        let mutation_data = mutations
            .into_iter()
            .map(TryInto::try_into)
            .collect::<Result<_, _>>()?;

        let results = repo
            .store_hg_changesets(changesets_data, mutation_data)
            .await?
            .into_iter()
            .map(move |r| {
                r.map(|(hg_cs_id, _bonsai_cs_id)| {
                    let hgid = HgId::from(hg_cs_id.into_nodehash());
                    UploadTokensResponse {
                        token: UploadToken::new_fake_token(AnyId::HgChangesetId(hgid), None),
                    }
                })
                .map_err(Error::from)
            });

        Ok(stream::iter(results).boxed())
    }
}

/// Upload list of bonsai changesets requested by the client
pub struct UploadBonsaiChangesetHandler;

#[async_trait]
impl EdenApiHandler for UploadBonsaiChangesetHandler {
    type QueryStringExtractor = UploadBonsaiChangesetQueryString;
    type Request = UploadBonsaiChangesetRequest;
    type Response = UploadTokensResponse;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::UploadBonsaiChangeset;
    const ENDPOINT: &'static str = "/upload/changeset/bonsai";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        let query = ectx.query();
        let bubble_id = query.bubble_id.map(BubbleId::new);
        let cs = request.changeset;
        let repo = &repo;
        let parents = stream::iter(cs.hg_parents)
            .then(|hgid| async move {
                repo.get_bonsai_from_hg(hgid.into())
                    .await?
                    .ok_or_else(|| anyhow!("Parent HgId {} is invalid", hgid))
            })
            .try_collect()
            .await?;
        let cs_id = repo
            .repo()
            .create_changeset(
                parents,
                CreateInfo {
                    author: cs.author,
                    author_date: DateTime::from_timestamp(cs.time, cs.tz)?.into(),
                    committer: None,
                    committer_date: None,
                    message: cs.message,
                    extra: cs.extra.into_iter().map(|e| (e.key, e.value)).collect(),
                    // TODO(rajshar): Need to allow passing git_extra_headers through Eden API as well.
                    git_extra_headers: None,
                },
                cs.file_changes
                    .into_iter()
                    .map(|(path, fc)| {
                        let create_change = to_create_change(fc, bubble_id)
                            .with_context(|| anyhow!("Parsing file changes for {}", path))?;
                        Ok((to_mononoke_path(path)?, create_change))
                    })
                    .collect::<anyhow::Result<_>>()?,
                match bubble_id {
                    Some(id) => Some(repo.open_bubble(id).await?),
                    None => None,
                }
                .as_ref(),
            )
            .await
            .with_context(|| anyhow!("When creating bonsai changeset"))?
            .id();

        Ok(stream::once(async move {
            Ok(UploadTokensResponse {
                token: UploadToken::new_fake_token(
                    AnyId::BonsaiChangesetId(cs_id.into()),
                    bubble_id.map(Into::into),
                ),
            })
        })
        .boxed())
    }
}

/// Get information about a snapshot changeset
pub struct FetchSnapshotHandler;

#[async_trait]
impl EdenApiHandler for FetchSnapshotHandler {
    type Request = FetchSnapshotRequest;
    type Response = FetchSnapshotResponse;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::FetchSnapshot;
    const ENDPOINT: &'static str = "/snapshot";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        let cs_id = ChangesetId::from(request.cs_id);
        let bubble_id = repo
            .ephemeral_store()
            .bubble_from_changeset(&cs_id)
            .await?
            .context("Snapshot not in a bubble")?;
        let labels = repo
            .ephemeral_store()
            .labels_from_bubble(&bubble_id)
            .await
            .context("Failed to fetch labels associated with the snapshot")?;
        let blobstore = repo.bubble_blobstore(Some(bubble_id)).await?;
        let cs = cs_id
            .load(repo.ctx(), &blobstore)
            .await
            .map_err(MononokeError::from)?
            .into_mut();
        let time = cs.author_date.timestamp_secs();
        let tz = cs.author_date.tz_offset_secs();
        let response = FetchSnapshotResponse {
            author: cs.author,
            time,
            tz,
            hg_parents: Parents::from_iter(
                stream::iter(
                    cs.parents
                        .into_iter()
                        .map(|cs_id| repo.get_hg_from_bonsai(cs_id)),
                )
                .buffered(2)
                .try_collect::<Vec<_>>()
                .await?
                .into_iter()
                .map(|id| id.into()),
            ),
            file_changes: cs
                .file_changes
                .into_iter()
                .map(|(path, fc)| {
                    Ok((
                        to_hg_path(&path)?,
                        match fc {
                            FileChange::Deletion => BonsaiFileChange::Deletion,
                            FileChange::UntrackedDeletion => BonsaiFileChange::UntrackedDeletion,
                            FileChange::Change(tc) => BonsaiFileChange::Change {
                                upload_token: UploadToken::new_fake_token(
                                    AnyId::AnyFileContentId(AnyFileContentId::ContentId(
                                        tc.content_id().into(),
                                    )),
                                    Some(bubble_id.into()),
                                ),
                                file_type: tc.file_type().into(),
                            },
                            FileChange::UntrackedChange(uc) => BonsaiFileChange::UntrackedChange {
                                upload_token: UploadToken::new_fake_token(
                                    AnyId::AnyFileContentId(AnyFileContentId::ContentId(
                                        uc.content_id().into(),
                                    )),
                                    Some(bubble_id.into()),
                                ),
                                file_type: uc.file_type().into(),
                            },
                        },
                    ))
                })
                .collect::<Result<_, Error>>()?,
            bubble_id: Some(bubble_id.into()),
            labels,
        };
        Ok(stream::once(async move { Ok(response) }).boxed())
    }
}

// Alter the properties of an existing snapshot
pub struct AlterSnapshotHandler;

#[async_trait]
impl EdenApiHandler for AlterSnapshotHandler {
    type Request = AlterSnapshotRequest;
    type Response = AlterSnapshotResponse;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::AlterSnapshot;
    const ENDPOINT: &'static str = "/snapshot/alter";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        let cs_id = ChangesetId::from(request.cs_id);
        let id = repo
            .ephemeral_store()
            .bubble_from_changeset(&cs_id)
            .await?
            .context("Snapshot does not exist or has already expired")?;
        let (label_addition, label_removal) = (
            !request.labels_to_add.is_empty(),
            !request.labels_to_remove.is_empty(),
        );
        if label_addition && label_removal {
            // Input has both labels to add and labels to remove, which is not allowed.
            Err(anyhow!(
                "Alter snapshot request cannot have labels_to_add and labels_to_remove both as non-empty"
            ))?
        } else if label_addition {
            // Input has labels to add, so let's add the input labels.
            repo.ephemeral_store()
                .add_bubble_labels(id, request.labels_to_add.clone())
                .await?;
        } else {
            // Input has labels to remove, or no labels as input at all. In either case,
            // we need to remove specific or all labels corresponding to the bubble.
            repo.ephemeral_store()
                .remove_bubble_labels(id, request.labels_to_remove.clone())
                .await?;
        }
        let current_labels = repo.ephemeral_store().labels_from_bubble(&id).await?;
        let response = AlterSnapshotResponse { current_labels };
        Ok(stream::once(async move { Ok(response) }).boxed())
    }
}

/// Creates an ephemeral bubble and return its id
pub struct EphemeralPrepareHandler;

#[async_trait]
impl EdenApiHandler for EphemeralPrepareHandler {
    type Request = EphemeralPrepareRequest;
    type Response = EphemeralPrepareResponse;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::EphemeralPrepare;
    const ENDPOINT: &'static str = "/ephemeral/prepare";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        Ok(stream::once(async move {
            Ok(EphemeralPrepareResponse {
                bubble_id: repo
                    .create_bubble(
                        request.custom_duration_secs.map(Duration::from_secs),
                        request.labels.unwrap_or_else(Vec::new),
                    )
                    .await?
                    .bubble_id()
                    .into(),
            })
        })
        .boxed())
    }
}

pub struct GraphHandlerV2;

#[async_trait]
impl EdenApiHandler for GraphHandlerV2 {
    type Request = CommitGraphRequest;
    type Response = CommitGraphEntry;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::CommitGraphV2;
    const ENDPOINT: &'static str = "/commit/graph_v2";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        let heads: Vec<_> = request
            .heads
            .into_iter()
            .map(|hg_id| HgChangesetId::new(HgNodeHash::from(hg_id)))
            .collect();
        let common: Vec<_> = request
            .common
            .into_iter()
            .map(|hg_id| HgChangesetId::new(HgNodeHash::from(hg_id)))
            .collect();

        if tunables()
            .enable_streaming_commit_graph_edenapi_endpoint()
            .unwrap_or_default()
        {
            // If all the requested heads are public, return stream.
            if heads.len() < PHASES_CHECK_LIMIT && repo.is_all_public(&heads).await? {
                let graph_stream = repo
                    .get_graph_mapping_stream(common, heads)
                    .await?
                    .err_into::<Error>()
                    .and_then(|(hgid, parents)| async move {
                        Ok(CommitGraphEntry {
                            hgid: HgId::from(hgid.into_nodehash()),
                            parents: parents
                                .into_iter()
                                .map(|p_hgid| HgId::from(p_hgid.into_nodehash()))
                                .collect(),
                            is_draft: Some(false),
                        })
                    })
                    .boxed();
                return Ok(graph_stream);
            }
        }

        let graph_entries = repo
            .get_graph_mapping(common, heads)
            .await?
            .into_iter()
            .map(|(hgid, (parents, is_draft))| {
                Ok(CommitGraphEntry {
                    hgid: HgId::from(hgid.into_nodehash()),
                    parents: parents
                        .into_iter()
                        .map(|p_hgid| HgId::from(p_hgid.into_nodehash()))
                        .collect(),
                    is_draft: Some(is_draft),
                })
            });
        Ok(stream::iter(graph_entries).boxed())
    }
}

pub struct GraphSegmentsHandler;

#[async_trait]
impl EdenApiHandler for GraphSegmentsHandler {
    type Request = CommitGraphSegmentsRequest;
    type Response = CommitGraphSegmentsEntry;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::CommitGraphSegments;
    const ENDPOINT: &'static str = "/commit/graph_segments";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        let heads: Vec<_> = request
            .heads
            .into_iter()
            .map(|hg_id| HgChangesetId::new(HgNodeHash::from(hg_id)))
            .collect();
        let common: Vec<_> = request
            .common
            .into_iter()
            .map(|hg_id| HgChangesetId::new(HgNodeHash::from(hg_id)))
            .collect();

        let graph_segment_entries =
            repo.graph_segments(common, heads)
                .await?
                .into_iter()
                .map(|segment| {
                    Ok(CommitGraphSegmentsEntry {
                        head: HgId::from(segment.head.into_nodehash()),
                        base: HgId::from(segment.base.into_nodehash()),
                        length: segment.length,
                        parents: segment
                            .parents
                            .into_iter()
                            .map(|parent| CommitGraphSegmentParent {
                                hgid: HgId::from(parent.hgid.into_nodehash()),
                                location: parent.location.map(|location| {
                                    location.map_descendant(|descendant| {
                                        HgId::from(descendant.into_nodehash())
                                    })
                                }),
                            })
                            .collect(),
                    })
                });
        Ok(stream::iter(graph_segment_entries).boxed())
    }
}

pub struct CommitMutationsHandler;

#[async_trait]
impl EdenApiHandler for CommitMutationsHandler {
    type Request = CommitMutationsRequest;
    type Response = CommitMutationsResponse;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::CommitMutations;
    const ENDPOINT: &'static str = "/commit/mutations";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        if !tunables().mutation_generate_for_draft().unwrap_or_default() {
            return Ok(stream::empty().boxed());
        }
        let commits = request
            .commits
            .into_iter()
            .map(|hg_id| HgChangesetId::new(HgNodeHash::from(hg_id)))
            .collect();

        let mutations = repo
            .fetch_mutations(commits)
            .await?
            .into_iter()
            .map(|mutation| {
                Ok(CommitMutationsResponse {
                    mutation: mutation.into(),
                })
            });

        Ok(stream::iter(mutations).boxed())
    }
}

pub struct CommitTranslateId;

#[async_trait]
impl EdenApiHandler for CommitTranslateId {
    type Request = CommitTranslateIdRequest;
    type Response = CommitTranslateIdResponse;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::CommitTranslateId;
    const ENDPOINT: &'static str = "/commit/translate_id";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let from_repo = match request.from_repo {
            Some(from_repo) => ectx.other_repo(from_repo).await?,
            None => ectx.repo(),
        };
        let from_repo = from_repo.repo();

        let to_repo = match request.to_repo {
            Some(to_repo) => ectx.other_repo(to_repo).await?,
            None => ectx.repo(),
        };
        let to_repo = to_repo.repo();

        let mut hg_ids = Vec::new();
        let mut bonsai_ids = Vec::new();
        let mut git_ids = Vec::new();
        let mut globalrevs = Vec::new();

        for commit in &request.commits {
            match commit {
                CommitId::Hg(hg_id) => hg_ids.push(HgChangesetId::from(*hg_id)),
                CommitId::Bonsai(bonsai_id) => bonsai_ids.push(ChangesetId::from(*bonsai_id)),
                CommitId::GitSha1(git_id) => git_ids.push(GitSha1::from(*git_id)),
                CommitId::Globalrev(globalrev) => globalrevs.push(Globalrev::new(*globalrev)),
            }
        }

        // Convert request types to intermediate bonsais using "from" repo.
        let (hg_bonsais, git_bonsais, globalrev_bonsais) = try_join!(
            from_repo.many_changeset_ids_from_hg(hg_ids),
            from_repo.many_changeset_ids_from_git_sha1(git_ids),
            from_repo.many_changeset_ids_from_globalrev(globalrevs),
        )?;

        // Mapping of request id to intermediate bonsai id of "from" repo.
        let mut input_to_bonsai: HashMap<CommitId, ChangesetId> = bonsai_ids
            .into_iter()
            .map(|id| (id.clone().into(), id))
            .chain(hg_bonsais.into_iter().map(|(hg, bs)| (hg.into(), bs)))
            .chain(git_bonsais.into_iter().map(|(g, bs)| (g.into(), bs)))
            .chain(globalrev_bonsais.into_iter().map(|(g, bs)| (g.into(), bs)))
            .collect();

        // Convert bonsai ids to that of "to" repo, if necessary.
        if from_repo.repoid() != to_repo.repoid() {
            input_to_bonsai = stream::iter(input_to_bonsai.into_iter())
                .then(|(id, bs)| {
                    let from_repo = from_repo.clone();
                    let to_repo = to_repo.clone();
                    async move {
                        (
                            id,
                            from_repo
                                .xrepo_commit_lookup(&to_repo, bs.clone(), None)
                                .await,
                        )
                    }
                })
                .filter_map(|(id, xctx)| async move {
                    match xctx {
                        // If there is no mapping, skip it.
                        // TODO: perhaps we should propagate an error
                        Ok(None) => None,
                        Ok(Some(xctx)) => Some(Ok((id, xctx.id()))),
                        Err(err) => Some(Err(err)),
                    }
                })
                .try_collect()
                .await?;
        }

        let all_bonsai_ids: Vec<_> = input_to_bonsai.values().cloned().collect();
        // Convert all bonsais to the target type
        let bonsai_to_target: HashMap<ChangesetId, CommitId> = match request.scheme {
            CommitIdScheme::Bonsai => all_bonsai_ids
                .into_iter()
                .map(|to| (to.clone(), to.into()))
                .collect(),
            CommitIdScheme::Hg => to_repo
                .many_changeset_hg_ids(all_bonsai_ids)
                .await?
                .into_iter()
                .map(|(id, hg_id)| (id, hg_id.into()))
                .collect(),
            CommitIdScheme::GitSha1 => to_repo
                .many_changeset_git_sha1s(all_bonsai_ids)
                .await?
                .into_iter()
                .map(|(id, git_sha1)| (id, git_sha1.into()))
                .collect(),
            CommitIdScheme::Globalrev => to_repo
                .many_changeset_globalrev_ids(all_bonsai_ids)
                .await?
                .into_iter()
                .map(|(id, globalrev)| (id, globalrev.into()))
                .collect(),
        };

        // Build the response based on the order within the request.
        let translations: Vec<_> = request
            .commits
            .into_iter()
            .filter_map(|commit| {
                let translated = bonsai_to_target.get(input_to_bonsai.get(&commit)?)?.clone();
                Some(CommitTranslateIdResponse { commit, translated })
            })
            .map(anyhow::Ok)
            .collect();

        Ok(stream::iter(translations).boxed())
    }
}


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

use anyhow::format_err;
use anyhow::Context;
use anyhow::Error;
use async_trait::async_trait;
use bytes::Bytes;
use context::PerfCounterType;
use edenapi_types::wire::ToWire;
use edenapi_types::AnyFileContentId;
use edenapi_types::AnyId;
use edenapi_types::Batch;
use edenapi_types::FileAttributes;
use edenapi_types::FileAuxData;
use edenapi_types::FileContent;
use edenapi_types::FileContentTokenMetadata;
use edenapi_types::FileEntry;
use edenapi_types::FileRequest;
use edenapi_types::FileResponse;
use edenapi_types::FileSpec;
use edenapi_types::ServerError;
use edenapi_types::UploadHgFilenodeRequest;
use edenapi_types::UploadToken;
use edenapi_types::UploadTokenMetadata;
use edenapi_types::UploadTokensResponse;
use ephemeral_blobstore::BubbleId;
use futures::stream;
use futures::Stream;
use futures::StreamExt;
use futures::TryStreamExt;
use gotham::state::FromState;
use gotham::state::State;
use gotham_derive::StateData;
use gotham_derive::StaticResponseExtender;
use gotham_ext::error::HttpError;
use gotham_ext::response::TryIntoResponse;
use hyper::Body;
use mercurial_types::HgFileNodeId;
use mercurial_types::HgNodeHash;
use mononoke_api_hg::HgDataContext;
use mononoke_api_hg::HgDataId;
use mononoke_api_hg::HgRepoContext;
use rate_limiting::Metric;
use serde::Deserialize;
use types::Key;

use super::handler::EdenApiContext;
use super::EdenApiHandler;
use super::EdenApiMethod;
use super::HandlerInfo;
use super::HandlerResult;
use crate::context::ServerContext;
use crate::errors::ErrorKind;
use crate::middleware::RequestContext;
use crate::utils::cbor_stream_filtered_errors;
use crate::utils::get_repo;

/// XXX: This number was chosen arbitrarily.
const MAX_CONCURRENT_FILE_FETCHES_PER_REQUEST: usize = 10;

const MAX_CONCURRENT_UPLOAD_FILENODES_PER_REQUEST: usize = 1000;

#[derive(Debug, Deserialize, StateData, StaticResponseExtender)]
pub struct UploadFileParams {
    repo: String,
    idtype: String,
    id: String,
}

#[derive(Debug, Deserialize, StateData, StaticResponseExtender)]
pub struct UploadFileQueryString {
    bubble_id: Option<NonZeroU64>,
    content_size: u64,
}

/// Fetch the content of the files requested by the client.
pub struct FilesHandler;

#[async_trait]
impl EdenApiHandler for FilesHandler {
    type Request = FileRequest;
    type Response = FileEntry;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::Files;
    const ENDPOINT: &'static str = "/files";

    fn sampling_rate(request: &Self::Request) -> NonZeroU64 {
        // Sample trivial requests
        if request.keys.len() + request.reqs.len() == 1 {
            nonzero_ext::nonzero!(100u64)
        } else {
            nonzero_ext::nonzero!(1u64)
        }
    }

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        let ctx = repo.ctx().clone();

        let len = request.keys.len() + request.reqs.len();
        let reqs = request
            .keys
            .into_iter()
            .map(|key| FileSpec {
                key,
                attrs: FileAttributes {
                    content: true,
                    aux_data: false,
                },
            })
            .chain(request.reqs.into_iter());
        ctx.perf_counters()
            .add_to_counter(PerfCounterType::EdenapiFiles, len as i64);
        let fetches = reqs.map(move |FileSpec { key, attrs }| fetch_file(repo.clone(), key, attrs));

        Ok(stream::iter(fetches)
            .buffer_unordered(MAX_CONCURRENT_FILE_FETCHES_PER_REQUEST)
            .inspect_ok(move |_| {
                ctx.session().bump_load(Metric::GetpackFiles, 1.0);
            })
            .boxed())
    }
}

/// Fetch the content of the files requested by the client.
pub struct Files2Handler;

#[async_trait]
impl EdenApiHandler for Files2Handler {
    type Request = FileRequest;
    type Response = FileResponse;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::Files2;
    const ENDPOINT: &'static str = "/files2";

    fn sampling_rate(_request: &Self::Request) -> NonZeroU64 {
        nonzero_ext::nonzero!(256u64)
    }

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        let ctx = repo.ctx().clone();

        let len = request.keys.len() + request.reqs.len();
        let reqs = request
            .keys
            .into_iter()
            .map(|key| FileSpec {
                key,
                attrs: FileAttributes {
                    content: true,
                    aux_data: false,
                },
            })
            .chain(request.reqs.into_iter());
        ctx.perf_counters()
            .add_to_counter(PerfCounterType::EdenapiFiles, len as i64);
        let fetches =
            reqs.map(move |FileSpec { key, attrs }| fetch_file_response(repo.clone(), key, attrs));

        Ok(stream::iter(fetches)
            .buffer_unordered(MAX_CONCURRENT_FILE_FETCHES_PER_REQUEST)
            .inspect(move |response| {
                if let Ok(result) = &response {
                    if result.result.is_ok() {
                        ctx.session().bump_load(Metric::GetpackFiles, 1.0);
                    }
                }
            })
            .boxed())
    }
}

async fn fetch_file_response(
    repo: HgRepoContext,
    key: Key,
    attrs: FileAttributes,
) -> Result<FileResponse, Error> {
    let result = fetch_file(repo, key.clone(), attrs)
        .await
        .map_err(|e| ServerError::generic(format!("{}", e)));
    Ok(FileResponse { key, result })
}

/// Fetch requested file for a single key.
/// Note that this function consumes the repo context in order
/// to construct a file context for the requested blob.
async fn fetch_file(
    repo: HgRepoContext,
    key: Key,
    attrs: FileAttributes,
) -> Result<FileEntry, Error> {
    let id = HgFileNodeId::from_node_hash(HgNodeHash::from(key.hgid));

    let ctx = id
        .context(repo)
        .await
        .with_context(|| ErrorKind::FileFetchFailed(key.clone()))?
        .with_context(|| ErrorKind::KeyDoesNotExist(key.clone()))?;

    let parents = ctx.hg_parents().into();
    let mut file = FileEntry::new(key.clone(), parents);

    if attrs.content {
        let (data, metadata) = ctx
            .content()
            .await
            .with_context(|| ErrorKind::FileFetchFailed(key.clone()))?;

        file = file.with_content(FileContent {
            hg_file_blob: data,
            metadata,
        });
    }

    if attrs.aux_data {
        let content_metadata = ctx
            .content_metadata()
            .await
            .with_context(|| ErrorKind::FileFetchFailed(key.clone()))?;

        file = file.with_aux_data(FileAuxData {
            total_size: content_metadata.total_size,
            content_id: content_metadata.content_id.into(),
            sha1: content_metadata.sha1.into(),
            sha256: content_metadata.sha256.into(),
            seeded_blake3: Some(content_metadata.seeded_blake3.into()),
        });
    }

    Ok(file)
}

/// Generate an upload token for alredy uploaded content
async fn generate_upload_token(
    _repo: HgRepoContext,
    id: AnyFileContentId,
    content_size: u64,
    bubble_id: Option<NonZeroU64>,
) -> Result<UploadToken, Error> {
    // At first, returns a fake token
    Ok(UploadToken::new_fake_token_with_metadata(
        AnyId::AnyFileContentId(id),
        bubble_id,
        UploadTokenMetadata::FileContentTokenMetadata(FileContentTokenMetadata { content_size }),
    ))
}

/// Upload content of a file
async fn store_file(
    repo: HgRepoContext,
    id: AnyFileContentId,
    data: impl Stream<Item = Result<Bytes, Error>> + Send,
    content_size: u64,
    bubble_id: Option<BubbleId>,
) -> Result<(), Error> {
    repo.store_file(id, content_size, data, bubble_id).await?;
    Ok(())
}

/// Upload content of a file requested by the client.
pub async fn upload_file(state: &mut State) -> Result<impl TryIntoResponse, HttpError> {
    let params = UploadFileParams::take_from(state);
    let query_string = UploadFileQueryString::take_from(state);

    state.put(HandlerInfo::new(&params.repo, EdenApiMethod::UploadFile));

    let rctx = RequestContext::borrow_from(state).clone();
    let sctx = ServerContext::borrow_from(state);

    let repo = get_repo(sctx, &rctx, &params.repo, None).await?;

    let id = AnyFileContentId::from_str(&format!("{}/{}", &params.idtype, &params.id))
        .map_err(HttpError::e400)?;

    let body = Body::take_from(state).map_err(Error::from);
    let content_size = query_string.content_size;

    store_file(
        repo.clone(),
        id.clone(),
        body,
        content_size,
        query_string.bubble_id.map(BubbleId::new),
    )
    .await
    .map_err(HttpError::e500)?;

    let token = generate_upload_token(repo, id, content_size, query_string.bubble_id)
        .await
        .map(|v| v.to_wire());

    Ok(cbor_stream_filtered_errors(super::monitor_request(
        state,
        stream::iter(vec![token]),
    )))
}

/// Store the content of a single HgFilenode
async fn store_hg_filenode(
    repo: HgRepoContext,
    item: UploadHgFilenodeRequest,
) -> Result<UploadTokensResponse, Error> {
    // TODO(liubovd): validate signature of the upload token (item.token) and
    // return 'ErrorKind::UploadHgFilenodeRequestInvalidToken' if it's invalid.
    // This will be added later, for now assume tokens are always valid.

    let node_id = item.data.node_id;
    let token = item.data.file_content_upload_token;

    let filenode: HgFileNodeId = HgFileNodeId::from_node_hash(HgNodeHash::from(node_id));

    let p1: Option<HgFileNodeId> = item
        .data
        .parents
        .p1()
        .cloned()
        .map(HgNodeHash::from)
        .map(HgFileNodeId::from_node_hash);

    let p2: Option<HgFileNodeId> = item
        .data
        .parents
        .p2()
        .cloned()
        .map(HgNodeHash::from)
        .map(HgFileNodeId::from_node_hash);

    let any_file_content_id = match token.data.id {
        AnyId::AnyFileContentId(id) => Some(id),
        _ => None,
    }
    .ok_or_else(|| {
        ErrorKind::UploadHgFilenodeRequestInvalidToken(
            node_id.clone(),
            "the provided token is not for file content".into(),
        )
    })?;

    let content_id = repo
        .convert_file_to_content_id(any_file_content_id, None)
        .await?
        .ok_or_else(|| format_err!("File from upload token should be present"))?;

    let content_size = match token.data.metadata {
        Some(UploadTokenMetadata::FileContentTokenMetadata(meta)) => meta.content_size,
        _ => repo.fetch_file_content_size(content_id, None).await?,
    };

    let metadata = Bytes::from(item.data.metadata);

    repo.store_hg_filenode(filenode, p1, p2, content_id, content_size, metadata)
        .await?;

    Ok(UploadTokensResponse {
        token: UploadToken::new_fake_token(AnyId::HgFilenodeId(node_id), None),
    })
}

/// Upload list of hg filenodes requested by the client (batch request).
pub struct UploadHgFilenodesHandler;

#[async_trait]
impl EdenApiHandler for UploadHgFilenodesHandler {
    type Request = Batch<UploadHgFilenodeRequest>;
    type Response = UploadTokensResponse;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::UploadHgFilenodes;
    const ENDPOINT: &'static str = "/upload/filenodes";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        let tokens = request
            .batch
            .into_iter()
            .map(move |item| store_hg_filenode(repo.clone(), item));
        Ok(stream::iter(tokens)
            .buffer_unordered(MAX_CONCURRENT_UPLOAD_FILENODES_PER_REQUEST)
            .boxed())
    }
}

/// Downloads a file given an upload token
pub struct DownloadFileHandler;

#[async_trait]
impl EdenApiHandler for DownloadFileHandler {
    type Request = UploadToken;
    type Response = Bytes;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::DownloadFile;
    const ENDPOINT: &'static str = "/download/file";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        let content = repo
            .download_file(request)
            .await?
            .context("File not found")?;
        Ok(content.boxed())
    }
}


/*
 * 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::num::NonZeroU64;

use async_trait::async_trait;
use edenapi_types::ToWire;
use futures::stream::BoxStream;
use gotham::extractor::PathExtractor;
use gotham::extractor::QueryStringExtractor;
use gotham_derive::StateData;
use gotham_derive::StaticResponseExtender;
use gotham_ext::error::HttpError;
use hyper::body::Body;
use mononoke_api::MononokeError;
use mononoke_api_hg::HgRepoContext;
use nonzero_ext::nonzero;
use serde::Deserialize;

use super::EdenApiMethod;
use crate::context::ServerContext;
use crate::middleware::RequestContext;
use crate::utils::get_repo;

pub trait PathExtractorWithRepo: PathExtractor<Body> + Send + Sync {
    fn repo(&self) -> &str;
}

#[derive(Debug, Deserialize, StateData, StaticResponseExtender)]
pub struct BasicPathExtractor {
    repo: String,
}

impl PathExtractorWithRepo for BasicPathExtractor {
    fn repo(&self) -> &str {
        &self.repo
    }
}

pub struct HandlerError(HttpError);

impl From<HandlerError> for HttpError {
    fn from(e: HandlerError) -> Self {
        e.0
    }
}

// Default errors to 500.
impl From<MononokeError> for HandlerError {
    fn from(e: MononokeError) -> Self {
        Self(HttpError::e500(e))
    }
}

impl From<anyhow::Error> for HandlerError {
    fn from(e: anyhow::Error) -> Self {
        Self(HttpError::e500(e))
    }
}

// Handlers can propagate HttpError for a specific response code.
impl From<HttpError> for HandlerError {
    fn from(e: HttpError) -> Self {
        Self(e)
    }
}

pub type HandlerResult<'a, Response> =
    Result<BoxStream<'a, anyhow::Result<Response>>, HandlerError>;

pub struct EdenApiContext<P, Q> {
    rctx: RequestContext,
    sctx: ServerContext,
    repo: HgRepoContext,
    path: P,
    query: Q,
}

impl<P, Q> EdenApiContext<P, Q> {
    pub fn new(
        rctx: RequestContext,
        sctx: ServerContext,
        repo: HgRepoContext,
        path: P,
        query: Q,
    ) -> Self {
        Self {
            rctx,
            sctx,
            repo,
            path,
            query,
        }
    }
    pub fn repo(&self) -> HgRepoContext {
        self.repo.clone()
    }

    #[allow(unused)]
    pub fn path(&self) -> &P {
        &self.path
    }

    pub fn query(&self) -> &Q {
        &self.query
    }

    /// Open an "other" repo (i.e. distinct from repo specified in URL path).
    pub async fn other_repo(&self, repo_name: impl AsRef<str>) -> Result<HgRepoContext, HttpError> {
        get_repo(&self.sctx, &self.rctx, repo_name, None).await
    }
}

#[async_trait]
pub trait EdenApiHandler: 'static {
    type PathExtractor: PathExtractorWithRepo = BasicPathExtractor;
    type QueryStringExtractor: QueryStringExtractor<Body> + Send + Sync =
        gotham::extractor::NoopQueryStringExtractor;
    type Request: ToWire + Send;
    type Response: ToWire + Send + 'static;

    const HTTP_METHOD: hyper::Method;
    const API_METHOD: EdenApiMethod;
    /// DON'T include the /:repo prefix.
    /// Example: "/ephemeral/prepare"
    const ENDPOINT: &'static str;

    fn sampling_rate(_request: &Self::Request) -> NonZeroU64 {
        nonzero!(1u64)
    }

    async fn handler(
        ctx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response>;
}


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

use anyhow::Context;
use anyhow::Error;
use async_trait::async_trait;
use cloned::cloned;
use edenapi_types::HistoryRequest;
use edenapi_types::HistoryResponseChunk;
use edenapi_types::WireHistoryEntry;
use futures::stream;
use futures::stream::BoxStream;
use futures::StreamExt;
use futures::TryStreamExt;
use mercurial_types::HgFileNodeId;
use mercurial_types::HgNodeHash;
use mononoke_api_hg::HgRepoContext;
use types::Key;

use super::handler::EdenApiContext;
use super::EdenApiHandler;
use super::EdenApiMethod;
use super::HandlerResult;
use crate::errors::ErrorKind;
use crate::utils::to_mpath;

type HistoryStream = BoxStream<'static, Result<WireHistoryEntry, Error>>;

/// XXX: This number was chosen arbitrarily.
const MAX_CONCURRENT_FETCHES_PER_REQUEST: usize = 10;

pub struct HistoryHandler;

#[async_trait]
impl EdenApiHandler for HistoryHandler {
    type Request = HistoryRequest;
    type Response = HistoryResponseChunk;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::History;
    const ENDPOINT: &'static str = "/history";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        let HistoryRequest { keys, length } = request;

        let fetches = keys.into_iter().map(move |key| {
            // Construct a Future that buffers the full history for this key.
            // This should be OK since the history entries are relatively
            // small, so unless the history is extremely long, the total
            // amount of buffered data should be reasonable.
            cloned!(repo);
            async move {
                let path = key.path.clone();
                let stream = fetch_history_for_key(repo, key, length).await?;
                let entries = stream.try_collect().await?;
                Ok(HistoryResponseChunk { path, entries })
            }
        });

        Ok(stream::iter(fetches)
            .buffer_unordered(MAX_CONCURRENT_FETCHES_PER_REQUEST)
            .boxed())
    }
}

async fn fetch_history_for_key(
    repo: HgRepoContext,
    key: Key,
    length: Option<u32>,
) -> Result<HistoryStream, Error> {
    let filenode_id = HgFileNodeId::new(HgNodeHash::from(key.hgid));
    let mpath = to_mpath(&key.path)?.context(ErrorKind::UnexpectedEmptyPath)?;

    let file = repo
        .file(filenode_id)
        .await
        .with_context(|| ErrorKind::FileFetchFailed(key.clone()))?
        .with_context(|| ErrorKind::KeyDoesNotExist(key.clone()))?;

    // Fetch the file's history and convert the entries into
    // the expected on-the-wire format.
    let history = file
        .history(mpath, length)
        .err_into::<Error>()
        .map_err(move |e| e.context(ErrorKind::HistoryFetchFailed(key.clone())))
        .and_then(|entry| async { WireHistoryEntry::try_from(entry) })
        .boxed();

    Ok(history)
}


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

use std::collections::HashMap;

use anyhow::Context;
use anyhow::Error;
use async_trait::async_trait;
use bookmarks_movement::BookmarkKindRestrictions;
use bytes::Bytes;
use edenapi_types::HgId;
use edenapi_types::LandStackRequest;
use edenapi_types::LandStackResponse;
use futures::stream;
use futures::StreamExt;
use hooks::PushAuthoredBy;
use mercurial_types::HgChangesetId;
use mercurial_types::HgNodeHash;
use mononoke_api_hg::HgRepoContext;

use super::handler::EdenApiContext;
use super::EdenApiHandler;
use super::EdenApiMethod;
use super::HandlerResult;
use crate::errors::ErrorKind;

/// Rebase a stack of commits onto a bookmark, and update the bookmark to the top of the newly-rebased stack.
pub struct LandStackHandler;

#[async_trait]
impl EdenApiHandler for LandStackHandler {
    type Request = LandStackRequest;
    type Response = LandStackResponse;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::LandStack;
    const ENDPOINT: &'static str = "/land";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        Ok(stream::once(land_stack(
            ectx.repo(),
            request.bookmark,
            request.head,
            request.base,
            request
                .pushvars
                .into_iter()
                .map(|p| (p.key, p.value.into()))
                .collect(),
        ))
        .boxed())
    }
}

async fn land_stack(
    repo: HgRepoContext,
    bookmark: String,
    head_hgid: HgId,
    base_hgid: HgId,
    pushvars: HashMap<String, Bytes>,
) -> Result<LandStackResponse, Error> {
    let repo = repo.repo();

    let head = HgChangesetId::new(HgNodeHash::from(head_hgid));
    let head = repo
        .changeset(head)
        .await
        .context("failed to resolve head")?
        .ok_or(ErrorKind::HgIdNotFound(head_hgid))?
        .id();

    let base = HgChangesetId::new(HgNodeHash::from(base_hgid));
    let base = repo
        .changeset(base)
        .await
        .context("failed to resolve base")?
        .ok_or(ErrorKind::HgIdNotFound(base_hgid))?
        .id();

    let pushrebase_outcome = repo
        .land_stack(
            bookmark,
            head,
            base,
            if pushvars.is_empty() {
                None
            } else {
                Some(&pushvars)
            },
            BookmarkKindRestrictions::AnyKind,
            PushAuthoredBy::User,
        )
        .await?;

    let new_head = pushrebase_outcome.head;
    let (old_ids, new_ids): (Vec<_>, Vec<_>) = pushrebase_outcome
        .rebased_changesets
        .into_iter()
        .map(|pair| (pair.id_old, pair.id_new))
        .unzip();

    let all_ids = std::iter::once(new_head)
        .chain(old_ids.iter().copied())
        .chain(new_ids.iter().copied())
        .collect();
    let all_hgids: HashMap<_, _> = repo
        .many_changeset_hg_ids(all_ids)
        .await?
        .into_iter()
        .collect();

    let new_head_hgid = all_hgids
        .get(&new_head)
        .ok_or(ErrorKind::BonsaiChangesetToHgIdError(new_head))
        .context("failed to fetch hgid for new head")?
        .into_nodehash()
        .into();

    let old_hgids: Result<Vec<_>, _> = old_ids
        .iter()
        .map(|id| {
            all_hgids
                .get(id)
                .ok_or(ErrorKind::BonsaiChangesetToHgIdError(*id))
                .context("failed to fetch hgids for old ids")
                .map(|id| id.into_nodehash().into())
        })
        .collect();

    let new_hgids: Result<Vec<_>, _> = new_ids
        .iter()
        .map(|id| {
            all_hgids
                .get(id)
                .ok_or(ErrorKind::BonsaiChangesetToHgIdError(*id))
                .context("failed to fetch hgids for new ids")
                .map(|id| id.into_nodehash().into())
        })
        .collect();

    let old_to_new_hgids = old_hgids?.into_iter().zip(new_hgids?.into_iter()).collect();

    Ok(LandStackResponse {
        new_head: new_head_hgid,
        old_to_new_hgids,
    })
}


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

use anyhow::bail;
use anyhow::Result;
use async_trait::async_trait;
use edenapi_types::AnyFileContentId;
use edenapi_types::AnyId;
use edenapi_types::Batch;
use edenapi_types::FileContentTokenMetadata;
use edenapi_types::IndexableId;
use edenapi_types::LookupRequest;
use edenapi_types::LookupResponse;
use edenapi_types::LookupResult;
use edenapi_types::UploadToken;
use edenapi_types::UploadTokenMetadata;
use ephemeral_blobstore::BubbleId;
use ephemeral_blobstore::StorageLocation;
use filestore::FilestoreConfigRef;
use futures::stream;
use futures::StreamExt;
use mercurial_types::HgChangesetId;
use mercurial_types::HgFileNodeId;
use mercurial_types::HgManifestId;
use mercurial_types::HgNodeHash;
use mononoke_api_hg::HgDataId;
use mononoke_api_hg::HgRepoContext;
use repo_blobstore::RepoBlobstoreRef;
use repo_identity::RepoIdentityRef;

use super::handler::EdenApiContext;
use super::EdenApiHandler;
use super::EdenApiMethod;
use super::HandlerResult;

const MAX_CONCURRENT_LOOKUPS_PER_REQUEST: usize = 10000;

enum Lookup {
    NotPresent,
    Present(Option<UploadTokenMetadata>),
}

impl From<bool> for Lookup {
    fn from(b: bool) -> Self {
        if b {
            Self::Present(None)
        } else {
            Self::NotPresent
        }
    }
}

async fn maybe_copy_file(
    repo: HgRepoContext,
    id: AnyFileContentId,
    bubble_id: Option<BubbleId>,
    copy_from_bubble_id: BubbleId,
) -> Result<Lookup> {
    Ok(if let Some(bubble_id) = bubble_id {
        let blob_repo = repo.repo().blob_repo();
        match repo
            .open_bubble(copy_from_bubble_id)
            .await?
            .copy_file_to_bubble(
                repo.ctx(),
                blob_repo.repo_identity().id(),
                blob_repo.repo_blobstore().clone(),
                bubble_id,
                *blob_repo.filestore_config(),
                id.into(),
            )
            .await?
        {
            None => Lookup::NotPresent,
            Some(data) => Lookup::Present(Some(
                FileContentTokenMetadata {
                    content_size: data.total_size,
                }
                .into(),
            )),
        }
    } else {
        Lookup::NotPresent
    })
}

async fn check_file(
    repo: HgRepoContext,
    id: AnyFileContentId,
    bubble_id: Option<BubbleId>,
    copy_from_bubble: Option<BubbleId>,
) -> Result<Lookup> {
    let content_id = repo.convert_file_to_content_id(id, bubble_id).await?;
    let lookup = if let Some(content_id) = content_id {
        // Reasons why check if content id is present:
        // 1. If content_id is provided, we haven't yet checked it is actually
        // in the blobstore
        // 2. Maybe alias was written to blobstore but the actual blob has not
        if repo.is_file_present(content_id, bubble_id).await? {
            Lookup::Present(Some(
                FileContentTokenMetadata {
                    content_size: repo.fetch_file_content_size(content_id, bubble_id).await?,
                }
                .into(),
            ))
        } else {
            Lookup::NotPresent
        }
    } else {
        Lookup::NotPresent
    };
    let lookup = match (lookup, copy_from_bubble) {
        (Lookup::NotPresent, Some(copy_bid)) => maybe_copy_file(repo, id, bubble_id, copy_bid)
            .await
            .unwrap_or(Lookup::NotPresent),
        (l, _) => l,
    };
    Ok(lookup)
}

/// Check if the item is present already and generate a token if it is.
/// Return None if the item has to be uploaded
/// Item can be of any type:
///     * file content id (canonical and sha1, sha256 aliases)
///     * hg filenode id
///     * hg tree id
///     * hg changeset id
async fn check_request_item(repo: HgRepoContext, item: LookupRequest) -> Result<LookupResponse> {
    let old_bubble_id = item.bubble_id;
    let bubble_id = old_bubble_id.map(BubbleId::new);
    if item.copy_from_bubble_id.is_some() && !matches!(item.id, AnyId::AnyFileContentId(_)) {
        bail!("copy_from_bubble_id is only supported with files")
    }
    if item.bubble_id.is_some()
        && matches!(
            item.id,
            AnyId::HgFilenodeId(_) | AnyId::HgTreeId(_) | AnyId::HgChangesetId(_)
        )
    {
        bail!("Hg derived data cannot be stored in bubbles")
    }
    let lookup = match item.id {
        AnyId::AnyFileContentId(id) => {
            check_file(
                repo,
                id,
                bubble_id,
                item.copy_from_bubble_id.map(BubbleId::new),
            )
            .await?
        }
        AnyId::BonsaiChangesetId(id) => repo
            .changeset_exists(
                id.into(),
                match bubble_id {
                    Some(id) => StorageLocation::Bubble(id),
                    None => StorageLocation::Persistent,
                },
            )
            .await?
            .into(),
        // Hg derived data does not exist on bubbles, let's fail fast
        AnyId::HgFilenodeId(id) => repo
            .filenode_exists(HgFileNodeId::from_node_hash(HgNodeHash::from(id)))
            .await?
            .into(),
        AnyId::HgTreeId(id) => repo
            .tree_exists(HgManifestId::new(HgNodeHash::from(id)))
            .await?
            .into(),
        AnyId::HgChangesetId(id) => repo
            .hg_changeset_exists(HgChangesetId::new(HgNodeHash::from(id)))
            .await?
            .into(),
    };
    let result = match lookup {
        Lookup::NotPresent => LookupResult::NotPresent(IndexableId {
            id: item.id,
            bubble_id: old_bubble_id,
        }),
        Lookup::Present(None) => {
            let token = UploadToken::new_fake_token(item.id.clone(), item.bubble_id);
            LookupResult::Present(token)
        }
        Lookup::Present(Some(metadata)) => {
            let token = UploadToken::new_fake_token_with_metadata(
                item.id.clone(),
                item.bubble_id,
                metadata,
            );
            LookupResult::Present(token)
        }
    };

    Ok(LookupResponse { result })
}

/// Process lookup (batched) request.
pub struct LookupHandler;

#[async_trait]
impl EdenApiHandler for LookupHandler {
    type Request = Batch<LookupRequest>;
    type Response = LookupResponse;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::Lookup;
    const ENDPOINT: &'static str = "/lookup";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        let tokens = request
            .batch
            .into_iter()
            .map(move |item| check_request_item(repo.clone(), item));

        Ok(stream::iter(tokens)
            .buffer_unordered(MAX_CONCURRENT_LOOKUPS_PER_REQUEST)
            .boxed())
    }
}


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

use anyhow::Context;
use anyhow::Error;
use edenapi_types::ToWire;
use futures::channel::oneshot;
use futures::stream::TryStreamExt;
use futures::FutureExt;
use futures::Stream;
use futures_stats::futures03::TimedFutureExt;
use gotham::handler::HandlerError as GothamHandlerError;
use gotham::handler::HandlerFuture;
use gotham::middleware::state::StateMiddleware;
use gotham::pipeline::new_pipeline;
use gotham::pipeline::single_pipeline;
use gotham::router::builder::build_router as gotham_build_router;
use gotham::router::builder::DefineSingleRoute;
use gotham::router::builder::DrawRoutes;
use gotham::router::builder::RouterBuilder;
use gotham::router::Router;
use gotham::state::FromState;
use gotham::state::State;
use gotham_derive::StateData;
use gotham_ext::content_encoding::ContentEncoding;
use gotham_ext::error::ErrorFormatter;
use gotham_ext::middleware::load::RequestLoad;
use gotham_ext::middleware::scuba::HttpScubaKey;
use gotham_ext::middleware::scuba::ScubaMiddlewareState;
use gotham_ext::middleware::MetadataState;
use gotham_ext::response::build_response;
use gotham_ext::response::encode_stream;
use gotham_ext::response::ResponseTryStreamExt;
use gotham_ext::response::StreamBody;
use gotham_ext::response::TryIntoResponse;
use gotham_ext::state_ext::StateExt;
use hyper::Body;
use hyper::Response;
use mime::Mime;
use serde::Deserialize;
use serde::Serialize;
use time_ext::DurationExt;
use tunables::tunables;

use crate::context::ServerContext;
use crate::middleware::request_dumper::RequestDumper;
use crate::middleware::RequestContext;
use crate::scuba::EdenApiScubaKey;
use crate::utils::cbor_mime;
use crate::utils::get_repo;
use crate::utils::monitor::Monitor;
use crate::utils::parse_wire_request;
use crate::utils::to_cbor_bytes;

mod blame;
mod bookmarks;
mod capabilities;
mod clone;
mod commit;
mod files;
mod handler;
mod history;
mod land;
mod lookup;
mod pull;
mod repos;
mod trees;

pub(crate) use handler::EdenApiHandler;
pub(crate) use handler::HandlerResult;
pub(crate) use handler::PathExtractorWithRepo;

use self::handler::EdenApiContext;

/// Enum identifying the EdenAPI method that each handler corresponds to.
/// Used to identify the handler for logging and stats collection.
#[derive(Copy, Clone)]
pub enum EdenApiMethod {
    Blame,
    Capabilities,
    Files,
    Files2,
    Lookup,
    UploadFile,
    UploadHgFilenodes,
    UploadTrees,
    UploadHgChangesets,
    UploadBonsaiChangeset,
    Trees,
    History,
    CommitLocationToHash,
    CommitHashToLocation,
    CommitRevlogData,
    CommitHashLookup,
    Clone,
    Bookmarks,
    SetBookmark,
    LandStack,
    PullFastForwardMaster,
    PullLazy,
    EphemeralPrepare,
    FetchSnapshot,
    AlterSnapshot,
    CommitGraphV2,
    CommitGraphSegments,
    DownloadFile,
    CommitMutations,
    CommitTranslateId,
}

impl fmt::Display for EdenApiMethod {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let name = match self {
            Self::Blame => "blame",
            Self::Capabilities => "capabilities",
            Self::Files => "files",
            Self::Files2 => "files2",
            Self::Trees => "trees",
            Self::History => "history",
            Self::CommitLocationToHash => "commit_location_to_hash",
            Self::CommitHashToLocation => "commit_hash_to_location",
            Self::CommitRevlogData => "commit_revlog_data",
            Self::CommitHashLookup => "commit_hash_lookup",
            Self::CommitGraphV2 => "commit_graph_v2",
            Self::CommitGraphSegments => "commit_graph_segments",
            Self::Clone => "clone",
            Self::Bookmarks => "bookmarks",
            Self::SetBookmark => "set_bookmark",
            Self::LandStack => "land_stack",
            Self::Lookup => "lookup",
            Self::UploadFile => "upload_file",
            Self::PullFastForwardMaster => "pull_fast_forward_master",
            Self::PullLazy => "pull_lazy",
            Self::UploadHgFilenodes => "upload_filenodes",
            Self::UploadTrees => "upload_trees",
            Self::UploadHgChangesets => "upload_hg_changesets",
            Self::UploadBonsaiChangeset => "upload_bonsai_changeset",
            Self::EphemeralPrepare => "ephemeral_prepare",
            Self::FetchSnapshot => "fetch_snapshot",
            Self::AlterSnapshot => "alter_snapshot",
            Self::DownloadFile => "download_file",
            Self::CommitMutations => "commit_mutations",
            Self::CommitTranslateId => "commit_translate_id",
        };
        write!(f, "{}", name)
    }
}

/// Information about the handler that served the request.
///
/// This should be inserted into the request's `State` by each handler. It will
/// typically be used by middlware for request logging and stats reporting.
#[derive(Default, StateData, Clone)]
pub struct HandlerInfo {
    pub repo: Option<String>,
    pub method: Option<EdenApiMethod>,
}

impl HandlerInfo {
    pub fn new(repo: impl ToString, method: EdenApiMethod) -> Self {
        Self {
            repo: Some(repo.to_string()),
            method: Some(method),
        }
    }
}

/// JSON representation of an error to send to the client.
#[derive(Clone, Serialize, Debug, Deserialize)]
struct JsonError {
    message: String,
    request_id: String,
}

struct JsonErrorFomatter;

impl ErrorFormatter for JsonErrorFomatter {
    type Body = Vec<u8>;

    fn format(&self, error: &Error, state: &State) -> Result<(Self::Body, Mime), Error> {
        let message = format!("{:#}", error);

        // Package the error message into a JSON response.
        let res = JsonError {
            message,
            request_id: state.short_request_id().to_string(),
        };

        let body = serde_json::to_vec(&res).context("Failed to serialize error")?;

        Ok((body, mime::APPLICATION_JSON))
    }
}

/// Macro to create a Gotham handler function from an async function.
///
/// The expected signature of the input function is:
/// ```rust,ignore
/// async fn handler(state: &mut State) -> Result<impl TryIntoResponse, HttpError>
/// ```
///
/// The resulting wrapped function will have the signaure:
/// ```rust,ignore
/// fn wrapped(mut state: State) -> Pin<Box<HandlerFuture>>
/// ```
macro_rules! define_handler {
    ($name:ident, $func:path) => {
        fn $name(mut state: State) -> Pin<Box<HandlerFuture>> {
            async move {
                let (future_stats, res) = $func(&mut state).timed().await;
                ScubaMiddlewareState::try_set_future_stats(&mut state, &future_stats);
                build_response(res, state, &JsonErrorFomatter)
            }
            .boxed()
        }
    };
}

define_handler!(repos_handler, repos::repos);
define_handler!(trees_handler, trees::trees);
define_handler!(capabilities_handler, capabilities::capabilities_handler);
define_handler!(commit_hash_to_location_handler, commit::hash_to_location);
define_handler!(commit_revlog_data_handler, commit::revlog_data);
define_handler!(clone_handler, clone::clone_data);
define_handler!(upload_file_handler, files::upload_file);
define_handler!(pull_fast_forward_master, pull::pull_fast_forward_master);
define_handler!(pull_lazy, pull::pull_lazy);

static HIGH_LOAD_SIGNAL: &str = "I_AM_OVERLOADED";
static ALIVE: &str = "I_AM_ALIVE";
static EXITING: &str = "EXITING";

// Used for monitoring VIP health
fn proxygen_health_handler(state: State) -> (State, &'static str) {
    if ServerContext::borrow_from(&state).will_exit() {
        (state, EXITING)
    } else {
        if let Some(request_load) = RequestLoad::try_borrow_from(&state) {
            let threshold = tunables::tunables()
                .edenapi_high_load_threshold()
                .unwrap_or_default();
            if threshold > 0 && request_load.0 > threshold {
                return (state, HIGH_LOAD_SIGNAL);
            }
        }
        (state, ALIVE)
    }
}

// Used for monitoring TW tasks
fn health_handler(state: State) -> (State, &'static str) {
    if ServerContext::borrow_from(&state).will_exit() {
        (state, EXITING)
    } else {
        (state, ALIVE)
    }
}

async fn handler_wrapper<Handler: EdenApiHandler>(
    mut state: State,
) -> Result<(State, Response<Body>), (State, GothamHandlerError)>
where
    <Handler as EdenApiHandler>::Request: std::fmt::Debug,
{
    let (future_stats, res) = async {
        let path = Handler::PathExtractor::take_from(&mut state);
        let query = Handler::QueryStringExtractor::take_from(&mut state);
        let content_encoding = ContentEncoding::from_state(&state);

        state.put(HandlerInfo::new(path.repo(), Handler::API_METHOD));

        let rctx = RequestContext::borrow_from(&state).clone();
        let sctx = ServerContext::borrow_from(&state).clone();

        let repo = get_repo(&sctx, &rctx, path.repo(), None).await?;
        let request = parse_wire_request::<<Handler::Request as ToWire>::Wire>(&mut state).await?;

        let sampling_rate = Handler::sampling_rate(&request);
        if sampling_rate.get() > 1 {
            ScubaMiddlewareState::try_set_sampling_rate(&mut state, sampling_rate);
        }

        if let Some(rd) = RequestDumper::try_borrow_mut_from(&mut state) {
            rd.add_request(&request);
        }

        let ectx = EdenApiContext::new(rctx, sctx, repo, path, query);

        match Handler::handler(ectx, request).await {
            Ok(responses) => Ok(encode_response_stream(
                monitor_request(&state, responses),
                content_encoding,
            )),
            Err(err) => Err(err.into()),
        }
    }
    .timed()
    .await;
    ScubaMiddlewareState::try_set_future_stats(&mut state, &future_stats);

    build_response(res, state, &JsonErrorFomatter)
}

pub fn monitor_request<S, T>(state: &State, stream: S) -> impl Stream<Item = T> + Send + 'static
where
    S: Stream<Item = T> + Send + 'static,
{
    let start = Instant::now();
    let ctx = RequestContext::borrow_from(state).ctx.clone();
    let (sender, receiver) = oneshot::channel::<()>();

    // EdenApi doesn't fill these in until the end of the request, so we need
    // to add them now.   A future improvement is to put these on the scuba
    // sample builder earlier on so we can clone it.
    let mut base_scuba = ctx.scuba().clone();
    base_scuba.add(
        HttpScubaKey::RequestId,
        state.short_request_id().to_string(),
    );

    if let Some(info) = state.try_borrow::<HandlerInfo>() {
        base_scuba.add_opt(EdenApiScubaKey::Repo, info.repo.clone());
        base_scuba.add_opt(EdenApiScubaKey::Method, info.method.map(|m| m.to_string()));
    }

    if let Some(metadata_state) = MetadataState::try_borrow_from(state) {
        let metadata = metadata_state.metadata();
        if let Some(ref address) = metadata.client_ip() {
            base_scuba.add(HttpScubaKey::ClientIp, address.to_string());
        }

        let identities = metadata.identities();
        let identities: Vec<_> = identities.iter().map(|i| i.to_string()).collect();
        base_scuba.add(HttpScubaKey::ClientIdentities, identities);
    }

    let reporting_loop = async move {
        loop {
            let interval = match tunables()
                .edenapi_request_monitor_interval()
                .unwrap_or_default()
                .try_into()
            {
                Ok(interval) if interval > 0 => interval,
                _ => break,
            };

            tokio::time::sleep(Duration::from_secs(interval)).await;
            let mut scuba = base_scuba.clone();
            ctx.perf_counters().insert_perf_counters(&mut scuba);
            scuba.log_with_msg(
                "Long running EdenAPI request",
                format!("{}", start.elapsed().as_micros_unchecked()),
            );
        }
    };

    tokio::task::spawn(async move {
        futures::pin_mut!(reporting_loop);
        let _ = futures::future::select(reporting_loop, receiver).await;
    });

    Monitor::new(stream, sender)
}

/// Encode a stream of EdenAPI responses into its final on-wire representation.
///
/// This involves converting each item to its wire format, CBOR serializing them, and then
/// optionally compressing the resulting byte stream based on the specified Content-Encoding.
pub fn encode_response_stream<S, T>(stream: S, encoding: ContentEncoding) -> impl TryIntoResponse
where
    S: Stream<Item = Result<T, Error>> + Send + 'static,
    T: ToWire + Send + 'static,
{
    let stream = stream.and_then(|item| async move { to_cbor_bytes(&item.to_wire()) });
    let stream = encode_stream(stream, encoding, None).capture_first_err();
    StreamBody::new(stream, cbor_mime())
}

// We use a struct here (rather than just a global function) just for the convenience
// of writing `Handlers::setup::<MyHandler>(route)`
// instead of `setup_handler::<MyHandler, _, _>(route)`, to make things clearer.
struct Handlers<C, P> {
    _phantom: (std::marker::PhantomData<C>, std::marker::PhantomData<P>),
}

impl<C, P> Handlers<C, P>
where
    C: gotham::pipeline::PipelineHandleChain<P> + Copy + Send + Sync + 'static,
    P: std::panic::RefUnwindSafe + Send + Sync + 'static,
{
    fn setup<Handler: EdenApiHandler>(route: &mut RouterBuilder<C, P>)
    where
        <Handler as EdenApiHandler>::Request: std::fmt::Debug,
    {
        route
            .request(
                vec![Handler::HTTP_METHOD],
                &format!("/:repo{}", Handler::ENDPOINT),
            )
            .with_path_extractor::<Handler::PathExtractor>()
            .with_query_string_extractor::<Handler::QueryStringExtractor>()
            .to_async(handler_wrapper::<Handler>);
    }
}

pub fn build_router(ctx: ServerContext) -> Router {
    let pipeline = new_pipeline().add(StateMiddleware::new(ctx)).build();
    let (chain, pipelines) = single_pipeline(pipeline);

    gotham_build_router(chain, pipelines, |route| {
        route.get("/health_check").to(health_handler);
        route.get("/repos").to(repos_handler);
        route
            .get("/proxygen/health_check")
            .to(proxygen_health_handler);
        Handlers::setup::<commit::EphemeralPrepareHandler>(route);
        Handlers::setup::<commit::UploadHgChangesetsHandler>(route);
        Handlers::setup::<commit::UploadBonsaiChangesetHandler>(route);
        Handlers::setup::<commit::LocationToHashHandler>(route);
        Handlers::setup::<commit::HashLookupHandler>(route);
        Handlers::setup::<files::FilesHandler>(route);
        Handlers::setup::<files::Files2Handler>(route);
        Handlers::setup::<files::UploadHgFilenodesHandler>(route);
        Handlers::setup::<bookmarks::BookmarksHandler>(route);
        Handlers::setup::<bookmarks::SetBookmarkHandler>(route);
        Handlers::setup::<land::LandStackHandler>(route);
        Handlers::setup::<history::HistoryHandler>(route);
        Handlers::setup::<lookup::LookupHandler>(route);
        Handlers::setup::<trees::UploadTreesHandler>(route);
        Handlers::setup::<commit::FetchSnapshotHandler>(route);
        Handlers::setup::<commit::AlterSnapshotHandler>(route);
        Handlers::setup::<commit::GraphHandlerV2>(route);
        Handlers::setup::<commit::GraphSegmentsHandler>(route);
        Handlers::setup::<files::DownloadFileHandler>(route);
        Handlers::setup::<commit::CommitMutationsHandler>(route);
        Handlers::setup::<commit::CommitTranslateId>(route);
        Handlers::setup::<blame::BlameHandler>(route);
        route.get("/:repo/health_check").to(health_handler);
        route
            .get("/:repo/capabilities")
            .with_path_extractor::<capabilities::CapabilitiesParams>()
            .to(capabilities_handler);
        route
            .post("/:repo/trees")
            .with_path_extractor::<trees::TreeParams>()
            .to(trees_handler);
        route
            .post("/:repo/commit/hash_to_location")
            .with_path_extractor::<commit::HashToLocationParams>()
            .to(commit_hash_to_location_handler);
        route
            .post("/:repo/commit/revlog_data")
            .with_path_extractor::<commit::RevlogDataParams>()
            .to(commit_revlog_data_handler);
        route
            .post("/:repo/clone")
            .with_path_extractor::<clone::CloneParams>()
            .to(clone_handler);
        route
            .post("/:repo/pull_fast_forward_master")
            .with_path_extractor::<pull::PullFastForwardParams>()
            .to(pull_fast_forward_master);
        route
            .post("/:repo/pull_lazy")
            .with_path_extractor::<pull::PullLazyParams>()
            .to(pull_lazy);
        route
            .put("/:repo/upload/file/:idtype/:id")
            .with_path_extractor::<files::UploadFileParams>()
            .with_query_string_extractor::<files::UploadFileQueryString>()
            .to(upload_file_handler);
    })
}


/*
 * 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 bytes::Bytes;
use edenapi_types::wire::pull::WirePullFastForwardRequest;
use edenapi_types::wire::pull::WirePullLazyRequest;
use edenapi_types::wire::ToWire;
use edenapi_types::wire::WireCloneData;
use edenapi_types::wire::WireIdMapEntry;
use gotham::state::FromState;
use gotham::state::State;
use gotham_derive::StateData;
use gotham_derive::StaticResponseExtender;
use gotham_ext::error::HttpError;
use gotham_ext::response::BytesBody;
use mercurial_types::HgChangesetId;
use serde::Deserialize;
use types::HgId;

use crate::context::ServerContext;
use crate::errors::MononokeErrorExt;
use crate::handlers::EdenApiMethod;
use crate::handlers::HandlerInfo;
use crate::middleware::request_dumper::RequestDumper;
use crate::middleware::RequestContext;
use crate::utils::cbor;
use crate::utils::get_repo;
use crate::utils::parse_wire_request;

#[derive(Debug, Deserialize, StateData, StaticResponseExtender)]
pub struct PullFastForwardParams {
    repo: String,
}

#[derive(Debug, Deserialize, StateData, StaticResponseExtender)]
pub struct PullLazyParams {
    repo: String,
}

pub async fn pull_lazy(state: &mut State) -> Result<BytesBody<Bytes>, HttpError> {
    let params = PullLazyParams::take_from(state);

    state.put(HandlerInfo::new(&params.repo, EdenApiMethod::PullLazy));
    let request = parse_wire_request::<WirePullLazyRequest>(state).await?;
    if let Some(rd) = RequestDumper::try_borrow_mut_from(state) {
        rd.add_request(&request);
    };

    let sctx = ServerContext::borrow_from(state);
    let rctx = RequestContext::borrow_from(state).clone();
    let hg_repo_ctx = get_repo(sctx, &rctx, &params.repo, None).await?;

    let common: Vec<HgChangesetId> = request.common.into_iter().map(Into::into).collect();
    let missing: Vec<HgChangesetId> = request.missing.into_iter().map(Into::into).collect();
    let clone_data = hg_repo_ctx
        .segmented_changelog_pull_data(common, missing)
        .await
        .map_err(|e| e.into_http_error("error getting segmented changelog data"))?;
    let idmap = clone_data
        .idmap
        .into_iter()
        .map(|(k, v)| WireIdMapEntry {
            dag_id: k.to_wire(),
            hg_id: HgId::from(v.into_nodehash()).to_wire(),
        })
        .collect();
    let wire_clone_data = WireCloneData {
        flat_segments: clone_data
            .flat_segments
            .segments
            .into_iter()
            .collect::<Vec<_>>()
            .to_wire(),
        idmap,
    };

    Ok(BytesBody::new(
        cbor::to_cbor_bytes(&wire_clone_data).map_err(HttpError::e500)?,
        cbor::cbor_mime(),
    ))
}

// TODO(quark): Remove this once callsites are migrated to pull_lazy
pub async fn pull_fast_forward_master(state: &mut State) -> Result<BytesBody<Bytes>, HttpError> {
    let params = PullFastForwardParams::take_from(state);

    state.put(HandlerInfo::new(
        &params.repo,
        EdenApiMethod::PullFastForwardMaster,
    ));
    let request = parse_wire_request::<WirePullFastForwardRequest>(state).await?;
    if let Some(rd) = RequestDumper::try_borrow_mut_from(state) {
        rd.add_request(&request);
    };

    let sctx = ServerContext::borrow_from(state);
    let rctx = RequestContext::borrow_from(state).clone();
    let hg_repo_ctx = get_repo(sctx, &rctx, &params.repo, None).await?;

    let old_master: HgChangesetId = request.old_master.into();
    let new_master: HgChangesetId = request.new_master.into();
    let clone_data = hg_repo_ctx
        .segmented_changelog_pull_data(vec![old_master], vec![new_master])
        .await
        .map_err(|e| e.into_http_error("error getting segmented changelog data"))?;
    let idmap = clone_data
        .idmap
        .into_iter()
        .map(|(k, v)| WireIdMapEntry {
            dag_id: k.to_wire(),
            hg_id: HgId::from(v.into_nodehash()).to_wire(),
        })
        .collect();
    let wire_clone_data = WireCloneData {
        flat_segments: clone_data
            .flat_segments
            .segments
            .into_iter()
            .collect::<Vec<_>>()
            .to_wire(),
        idmap,
    };

    Ok(BytesBody::new(
        cbor::to_cbor_bytes(&wire_clone_data).map_err(HttpError::e500)?,
        cbor::cbor_mime(),
    ))
}


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

use anyhow::Context;
use bytes::Bytes;
use gotham::state::FromState;
use gotham::state::State;
use gotham_ext::error::HttpError;
use gotham_ext::response::BytesBody;
use serde::Serialize;

use crate::context::ServerContext;
use crate::errors::ErrorKind;

#[derive(Clone, Serialize, Debug)]
struct ReposResponse {
    repos: Vec<String>,
}

pub async fn repos(state: &mut State) -> Result<BytesBody<Bytes>, HttpError> {
    let sctx = ServerContext::borrow_from(state);
    let mononoke = sctx.mononoke_api();

    let repos = mononoke.repo_names().collect::<Vec<_>>();
    let response = ReposResponse { repos };
    let bytes: Bytes = serde_json::to_vec(&response)
        .context(ErrorKind::SerializationFailed)
        .map_err(HttpError::e500)?
        .into();

    Ok(BytesBody::new(bytes, mime::APPLICATION_JSON))
}


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

use anyhow::Context;
use anyhow::Error;
use async_trait::async_trait;
use bytes::Bytes;
use context::PerfCounterType;
use edenapi_types::wire::WireTreeRequest;
use edenapi_types::AnyId;
use edenapi_types::Batch;
use edenapi_types::EdenApiServerError;
use edenapi_types::FileMetadata;
use edenapi_types::TreeChildEntry;
use edenapi_types::TreeEntry;
use edenapi_types::TreeRequest;
use edenapi_types::UploadToken;
use edenapi_types::UploadTreeRequest;
use edenapi_types::UploadTreeResponse;
use futures::stream;
use futures::Future;
use futures::FutureExt;
use futures::Stream;
use futures::StreamExt;
use futures::TryStreamExt;
use gotham::state::FromState;
use gotham::state::State;
use gotham_derive::StateData;
use gotham_derive::StaticResponseExtender;
use gotham_ext::error::HttpError;
use gotham_ext::middleware::scuba::ScubaMiddlewareState;
use gotham_ext::response::TryIntoResponse;
use manifest::Entry;
use manifest::Manifest;
use mercurial_types::FileType;
use mercurial_types::HgFileNodeId;
use mercurial_types::HgManifestId;
use mercurial_types::HgNodeHash;
use mononoke_api_hg::HgDataContext;
use mononoke_api_hg::HgDataId;
use mononoke_api_hg::HgRepoContext;
use mononoke_api_hg::HgTreeContext;
use rate_limiting::Metric;
use serde::Deserialize;
use tunables::tunables;
use types::Key;
use types::RepoPathBuf;

use super::handler::EdenApiContext;
use super::EdenApiHandler;
use super::EdenApiMethod;
use super::HandlerInfo;
use super::HandlerResult;
use crate::context::ServerContext;
use crate::errors::ErrorKind;
use crate::middleware::request_dumper::RequestDumper;
use crate::middleware::RequestContext;
use crate::utils::custom_cbor_stream;
use crate::utils::get_repo;
use crate::utils::parse_wire_request;

/// XXX: This number was chosen arbitrarily.
const MAX_CONCURRENT_TREE_FETCHES_PER_REQUEST: usize = 10;
const MAX_CONCURRENT_METADATA_FETCHES_PER_TREE_FETCH: usize = 100;
const MAX_CONCURRENT_UPLOAD_TREES_PER_REQUEST: usize = 100;

#[derive(Debug, Deserialize, StateData, StaticResponseExtender)]
pub struct TreeParams {
    repo: String,
}

/// Fetch the tree nodes requested by the client.
pub async fn trees(state: &mut State) -> Result<impl TryIntoResponse, HttpError> {
    let params = TreeParams::take_from(state);

    state.put(HandlerInfo::new(&params.repo, EdenApiMethod::Trees));

    let rctx = RequestContext::borrow_from(state).clone();
    let sctx = ServerContext::borrow_from(state);

    let repo = get_repo(sctx, &rctx, &params.repo, Metric::TotalManifests).await?;
    let request = parse_wire_request::<WireTreeRequest>(state).await?;
    if let Some(rd) = RequestDumper::try_borrow_mut_from(state) {
        rd.add_request(&request);
    };
    repo.ctx()
        .perf_counters()
        .add_to_counter(PerfCounterType::EdenapiTrees, request.keys.len() as i64);

    ScubaMiddlewareState::try_set_sampling_rate(state, nonzero_ext::nonzero!(256_u64));

    Ok(custom_cbor_stream(
        super::monitor_request(state, fetch_all_trees(repo, request)),
        |tree_entry| tree_entry.as_ref().err(),
    ))
}

/// Fetch trees for all of the requested keys concurrently.
fn fetch_all_trees(
    repo: HgRepoContext,
    request: TreeRequest,
) -> impl Stream<Item = Result<TreeEntry, EdenApiServerError>> {
    let ctx = repo.ctx().clone();

    let fetch_metadata = request.attributes.child_metadata;
    let fetches = request.keys.into_iter().map(move |key| {
        fetch_tree(repo.clone(), key.clone(), fetch_metadata)
            .map(|r| r.map_err(|e| EdenApiServerError::with_key(key, e)))
    });

    stream::iter(fetches)
        .buffer_unordered(MAX_CONCURRENT_TREE_FETCHES_PER_REQUEST)
        .inspect_ok(move |_| {
            ctx.session().bump_load(Metric::TotalManifests, 1.0);
        })
}

/// Fetch requested tree for a single key.
/// Note that this function consumes the repo context in order
/// to construct a tree context for the requested blob.
async fn fetch_tree(
    repo: HgRepoContext,
    key: Key,
    fetch_metadata: bool,
) -> Result<TreeEntry, Error> {
    let id = HgManifestId::from_node_hash(HgNodeHash::from(key.hgid));

    let ctx = id
        .context(repo.clone())
        .await
        .with_context(|| ErrorKind::TreeFetchFailed(key.clone()))?
        .with_context(|| ErrorKind::KeyDoesNotExist(key.clone()))?;

    let (data, _) = ctx
        .content()
        .await
        .with_context(|| ErrorKind::TreeFetchFailed(key.clone()))?;
    let parents = ctx.hg_parents().into();

    let mut entry = TreeEntry::new(key.clone(), data, parents);

    if fetch_metadata {
        if let Some(entries) = fetch_child_metadata_entries(&repo, &ctx).await? {
            let children: Vec<Result<TreeChildEntry, EdenApiServerError>> = entries
                .buffer_unordered(MAX_CONCURRENT_METADATA_FETCHES_PER_TREE_FETCH)
                .map(|r| r.map_err(|e| EdenApiServerError::with_key(key.clone(), e)))
                .collect()
                .await;

            entry.with_children(Some(children));
        }
    }

    Ok(entry)
}

async fn fetch_child_metadata_entries<'a>(
    repo: &'a HgRepoContext,
    ctx: &'a HgTreeContext,
) -> Result<
    Option<impl Stream<Item = impl Future<Output = Result<TreeChildEntry, Error>> + 'a> + 'a>,
    Error,
> {
    let manifest = ctx.clone().into_blob_manifest()?;
    if let Some(limit) = tunables().edenapi_large_tree_metadata_limit() {
        if manifest.content().files.len() > limit as usize {
            return Ok(None);
        }
    }
    let entries = manifest.list().collect::<Vec<_>>();

    Ok(Some(
        stream::iter(entries)
            // .entries iterator is not `Send`
            .map({
                move |(name, entry)| async move {
                    let name = RepoPathBuf::from_string(name.to_string())?;
                    Ok(match entry {
                        Entry::Leaf((file_type, child_id)) => {
                            let child_key = Key::new(name, child_id.into_nodehash().into());
                            fetch_child_file_metadata(repo, file_type, child_key.clone()).await?
                        }
                        Entry::Tree(child_id) => TreeChildEntry::new_directory_entry(Key::new(
                            name,
                            child_id.into_nodehash().into(),
                        )),
                    })
                }
            }),
    ))
}

async fn fetch_child_file_metadata(
    repo: &HgRepoContext,
    file_type: FileType,
    child_key: Key,
) -> Result<TreeChildEntry, Error> {
    let metadata = repo
        .file(HgFileNodeId::new(child_key.hgid.into()))
        .await?
        .ok_or_else(|| ErrorKind::FileFetchFailed(child_key.clone()))?
        .content_metadata()
        .await?;
    Ok(TreeChildEntry::new_file_entry(
        child_key,
        FileMetadata {
            file_type: Some(file_type.into()),
            size: Some(metadata.total_size),
            content_sha1: Some(metadata.sha1.into()),
            content_sha256: Some(metadata.sha256.into()),
            content_id: Some(metadata.content_id.into()),
            content_seeded_blake3: Some(metadata.seeded_blake3.into()),
            ..Default::default()
        },
    ))
}

/// Store the content of a single tree
async fn store_tree(
    repo: HgRepoContext,
    item: UploadTreeRequest,
) -> Result<UploadTreeResponse, Error> {
    let upload_node_id = HgNodeHash::from(item.entry.node_id);
    let contents = item.entry.data;
    let p1 = item.entry.parents.p1().cloned().map(HgNodeHash::from);
    let p2 = item.entry.parents.p2().cloned().map(HgNodeHash::from);
    repo.store_tree(upload_node_id, p1, p2, Bytes::from(contents))
        .await?;
    Ok(UploadTreeResponse {
        token: UploadToken::new_fake_token(AnyId::HgTreeId(item.entry.node_id), None),
    })
}

/// Upload list of trees requested by the client (batch request).
pub struct UploadTreesHandler;

#[async_trait]
impl EdenApiHandler for UploadTreesHandler {
    type Request = Batch<UploadTreeRequest>;
    type Response = UploadTreeResponse;

    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
    const API_METHOD: EdenApiMethod = EdenApiMethod::UploadTrees;
    const ENDPOINT: &'static str = "/upload/trees";

    async fn handler(
        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
        request: Self::Request,
    ) -> HandlerResult<'async_trait, Self::Response> {
        let repo = ectx.repo();
        let tokens = request
            .batch
            .into_iter()
            .map(move |item| store_tree(repo.clone(), item));

        Ok(stream::iter(tokens)
            .buffer_unordered(MAX_CONCURRENT_UPLOAD_TREES_PER_REQUEST)
            .boxed())
    }
}


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

pub mod ods;
pub mod request_context;
pub mod request_dumper;

pub use self::ods::OdsMiddleware;
pub use self::request_context::RequestContext;
pub use self::request_context::RequestContextMiddleware;
pub use self::request_dumper::RequestDumperMiddleware;