// Copyright 2024 The NativeLink Authors. All rights reserved.

//

// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");

// you may not use this file except in compliance with the License.

// You may obtain a copy of the License at

//

//    See LICENSE file for details

//

// Unless required by applicable law or agreed to in writing, software

// distributed under the License is distributed on an "AS IS" BASIS,

// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

// See the License for the specific language governing permissions and

// limitations under the License.

use core::convert::Into;

use core::pin::Pin;

use core::time::Duration;

use std::collections::{HashMap, VecDeque};

use bytes::Bytes;

use futures::stream::{FuturesUnordered, Stream};

use futures::{StreamExt, TryStreamExt};

use nativelink_config::cas_server::{CasStoreConfig, WithInstanceName};

use nativelink_error::{Code, Error, ResultExt, error_if, make_err, make_input_err};

use nativelink_proto::build::bazel::remote::execution::v2::content_addressable_storage_server::{

    ContentAddressableStorage, ContentAddressableStorageServer as Server,

};

use nativelink_proto::build::bazel::remote::execution::v2::{

    BatchReadBlobsRequest, BatchReadBlobsResponse, BatchUpdateBlobsRequest,

    BatchUpdateBlobsResponse, Directory, FindMissingBlobsRequest, FindMissingBlobsResponse,

    GetTreeRequest, GetTreeResponse, batch_read_blobs_response, batch_update_blobs_response,

    compressor,

};

use nativelink_proto::google::rpc::Status as GrpcStatus;

use nativelink_store::ac_utils::get_and_decode_digest;

use nativelink_store::grpc_store::GrpcStore;

use nativelink_store::store_manager::StoreManager;

use nativelink_util::common::DigestInfo;

use nativelink_util::digest_hasher::make_ctx_for_hash_func;

use nativelink_util::store_trait::{Store, StoreLike};

use opentelemetry::context::FutureExt;

use tonic::{Request, Response, Status};

use tracing::{Instrument, Level, debug, error_span, instrument};

#[derive(Debug)]

pub struct CasServer {

    stores: HashMap<String, Store>,

}

type GetTreeStream = Pin<Box<dyn Stream<Item = Result<GetTreeResponse, Status>> + Send + 'static>>;

/// Per-blob deadline applied inside `BatchReadBlobs` / `BatchUpdateBlobs`.

const BATCH_PER_BLOB_TIMEOUT: Duration = Duration::from_secs(30);

impl CasServer {

    pub fn new(

        configs: &[WithInstanceName<CasStoreConfig>],

        store_manager: &StoreManager,

    ) -> Result<Self, Error> {

        let mut stores = HashMap::with_capacity(configs.len());

        for config in configs {

            let store = store_manager.get_store(&config.cas_store).ok_or_else(|| 
{0

                make_input_err!("'cas_store': '{}' does not exist", config.cas_store)

            })?;

            stores.insert(config.instance_name.clone(), store);

        }

        Ok(Self { stores })

    }

    pub fn into_service(self) -> Server<Self> {

        Server::new(self)

    }

    async fn inner_find_missing_blobs(

        &self,

        request: FindMissingBlobsRequest,

    ) -> Result<Response<FindMissingBlobsResponse>, Error> {

        let instance_name = &request.instance_name;

        let store = self

            .stores

            .get(instance_name)

            .err_tip(|| 
format!0
("'instance_name' not configured for '{instance_name}'"))
?0

            .clone();

        let mut requested_blobs = Vec::with_capacity(request.blob_digests.len());

        for digest in 
&request.blob_digests4
 {

            requested_blobs.
push6
(DigestInfo::try_from(digest.clone())
?1
.
into6
());

        }

        let sizes = store

            .has_many(&requested_blobs)

            .await

            .err_tip(|| "In find_missing_blobs")
?0
;

        let missing_blob_digests = sizes

            .into_iter()

            .zip(request.blob_digests)

            .
filter_map3
(|(maybe_size, digest)| maybe_size.map_or_else(|| Some(
digest2
), |_| None))

            .collect();

        Ok(Response::new(FindMissingBlobsResponse {

            missing_blob_digests,

        }))

    }

    async fn inner_batch_update_blobs(

        &self,

        request: BatchUpdateBlobsRequest,

    ) -> Result<Response<BatchUpdateBlobsResponse>, Error> {

        let instance_name = &request.instance_name;

        let store = self

            .stores

            .get(instance_name)

            .err_tip(|| 
format!0
("'instance_name' not configured for '{instance_name}'"))
?0

            .clone();

        // If we are a GrpcStore we shortcut here, as this is a special store.

        // Note: We don't know the digests here, so we try perform a very shallow

        // check to see if it's a grpc store.

        if let Some(
grpc_store0
) = store.downcast_ref::<GrpcStore>(None) {

            return grpc_store.batch_update_blobs(Request::new(request)).await;

        }

        let store_ref = &store;

        let update_futures: FuturesUnordered<_> = request

            .requests

            .into_iter()

            .
map3
(|request| async move {

                let digest = request

                    .digest

                    .clone()

                    .err_tip(|| "Digest not found in request")
?0
;

                let request_data = request.data;

                let digest_info = DigestInfo::try_from(digest.clone())
?0
;

                let size_bytes = usize::try_from(digest_info.size_bytes())

                    .err_tip(|| "Digest size_bytes was not convertible to usize")
?0
;

                error_if!(

                    size_bytes != request_data.len(),

                    "Digest for upload had mismatching sizes, digest said {} data  said {}",

                    size_bytes,

                    request_data.len()

                );

                // Apply a per-blob deadline so one slow upload does not

                // make the whole batch hit the client's overall deadline.

                let result = match tokio::time::timeout(

                    BATCH_PER_BLOB_TIMEOUT,

                    store_ref.update_oneshot(digest_info, request_data),

                )

                .await

                {

                    Ok(r) => r.err_tip(|| "Error writing to store"),

                    Err(_elapsed) => Err(make_err!(

                        Code::DeadlineExceeded,

                        "BatchUpdateBlobs per-blob timeout ({} s) elapsed for digest {}",

                        BATCH_PER_BLOB_TIMEOUT.as_secs(),

                        digest_info,

                    )),

                };

                Ok::<_, Error>(batch_update_blobs_response::Response {

                    digest: Some(digest),

                    status: Some(result.map_or_else(Into::into, |()| 
GrpcStatus::default3
())),

                })

            })

            .collect();

        let responses = update_futures

            .try_collect::<Vec<batch_update_blobs_response::Response>>()

            .await
?0
;

        Ok(Response::new(BatchUpdateBlobsResponse { responses }))

    }

    async fn inner_batch_read_blobs(

        &self,

        request: BatchReadBlobsRequest,

    ) -> Result<Response<BatchReadBlobsResponse>, Error> {

        let instance_name = &request.instance_name;

        let store = self

            .stores

            .get(instance_name)

            .err_tip(|| 
format!0
("'instance_name' not configured for '{instance_name}'"))
?0

            .clone();

        // If we are a GrpcStore we shortcut here, as this is a special store.

        // Note: We don't know the digests here, so we try perform a very shallow

        // check to see if it's a grpc store.

        if let Some(
grpc_store0
) = store.downcast_ref::<GrpcStore>(None) {

            return grpc_store.batch_read_blobs(Request::new(request)).await;

        }

        let store_ref = &store;

        let read_futures: FuturesUnordered<_> = request

            .digests

            .into_iter()

            .
map2
(|digest| async move {

                let digest_copy = DigestInfo::try_from(digest.clone())
?0
;

                // TODO(palfrey) There is a security risk here of someone taking all the memory on the instance.

                // Apply a per-blob deadline so one slow read does not

                // make the whole batch hit the client's overall deadline.

                let result = match tokio::time::timeout(

                    BATCH_PER_BLOB_TIMEOUT,

                    store_ref.get_part_unchunked(digest_copy, 0, None),

                )

                .await

                {

                    Ok(r) => r.err_tip(|| "Error reading from store"),

                    Err(_elapsed) => Err(make_err!(

                        Code::DeadlineExceeded,

                        "BatchReadBlobs per-blob timeout ({} s) elapsed for digest {}",

                        BATCH_PER_BLOB_TIMEOUT.as_secs(),

                        digest_copy,

                    )),

                };

                let (status, data) = result.map_or_else(

                    |mut e| {

                        if e.code == Code::NotFound {

                            // Trim the error code. Not Found is quite common and we don't want to send a large

                            // error (debug) message for something that is common. We resize to just the last

                            // message as it will be the most relevant.

                            e.messages.resize_with(1, String::new);

                        }

                        (e.into(), Bytes::new())

                    },

                    |v| (GrpcStatus::default(), v),

                );

                Ok::<_, Error>(batch_read_blobs_response::Response {

                    status: Some(status),

                    digest: Some(digest),

                    compressor: compressor::Value::Identity.into(),

                    data,

                })

            })

            .collect();

        let responses = read_futures

            .try_collect::<Vec<batch_read_blobs_response::Response>>()

            .await
?0
;

        Ok(Response::new(BatchReadBlobsResponse { responses }))

    }

    async fn inner_get_tree(

        &self,

        request: GetTreeRequest,

    ) -> Result<impl Stream<Item = Result<GetTreeResponse, Status>> + Send + use<>, Error> {

        let instance_name = &request.instance_name;

        let store = self

            .stores

            .get(instance_name)

            .err_tip(|| 
format!0
("'instance_name' not configured for '{instance_name}'"))
?0

            .clone();

        // If we are a GrpcStore we shortcut here, as this is a special store.

        // Note: We don't know the digests here, so we try perform a very shallow

        // check to see if it's a grpc store.

        if let Some(
grpc_store0
) = store.downcast_ref::<GrpcStore>(None) {

            let stream = grpc_store

                .get_tree(Request::new(request))

                .await?

                .into_inner();

            return Ok(stream.left_stream());

        }

        let root_digest: DigestInfo = request

            .root_digest

            .err_tip(|| "Expected root_digest to exist in GetTreeRequest")
?0

            .try_into()

            .err_tip(|| "In GetTreeRequest::root_digest")
?0
;

        let mut deque: VecDeque<DigestInfo> = VecDeque::new();

        let mut directories: Vec<Directory> = Vec::new();

        // `page_token` will return the `{hash_str}-{size_bytes}` of the current request's first directory digest.

        let page_token_digest = if request.page_token.is_empty() {

            root_digest

        } else {

            let mut page_token_parts = request.page_token.split('-');

            DigestInfo::try_new(

                page_token_parts

                    .next()

                    .err_tip(|| "Failed to parse `hash_str` in `page_token`")
?0
,

                page_token_parts

                    .next()

                    .err_tip(|| "Failed to parse `size_bytes` in `page_token`")
?0

                    .parse::<i64>()

                    .err_tip(|| "Failed to parse `size_bytes` as i64")
?0
,

            )

            .err_tip(|| "Failed to parse `page_token` as `Digest` in `GetTreeRequest`")
?0

        };

        let page_size = request.page_size;

        // If `page_size` is 0, paging is not necessary.

        let mut page_token_matched = page_size == 0;

        deque.push_back(root_digest);

        while !deque.is_empty() {

            let digest: DigestInfo = deque.pop_front().err_tip(|| "In VecDeque::pop_front")
?0
;

            let directory = get_and_decode_digest::<Directory>(&store, digest.into())

                .await

                .err_tip(|| "Converting digest to Directory")
?0
;

            if digest == page_token_digest {

                page_token_matched = true;

            }

            for directory in 
&directory.directories26
 {

                let digest: DigestInfo = directory

                    .digest

                    .clone()

                    .err_tip(|| "Expected Digest to exist in Directory::directories::digest")
?0

                    .try_into()

                    .err_tip(|| "In Directory::file::digest")
?0
;

                deque.push_back(digest);

            }

            let page_size_usize = usize::try_from(page_size).unwrap_or(usize::MAX);

            if page_token_matched {

                directories.push(directory);

                if directories.len() == page_size_usize {

                    break;

                }

            }

        }

        // `next_page_token` will return the `{hash_str}:{size_bytes}` of the next request's first directory digest.

        // It will be an empty string when it reached the end of the directory tree.

        let next_page_token: String = deque

            .front()

            .map_or_else(String::new, |value| 
format!3
("{value}"));

        Ok(futures::stream::once(async {

            Ok(GetTreeResponse {

                directories,

                next_page_token,

            })

        })

        .right_stream())

    }

}

#[tonic::async_trait]

impl ContentAddressableStorage for CasServer {

    type GetTreeStream = GetTreeStream;

    #[instrument(

        err,

        ret(level = Level::DEBUG),

        level = Level::ERROR,

        skip_all,

        fields(

            // Mostly to skip request.blob_digests which is sometimes enormous

            request.instance_name = ?grpc_request.get_ref().instance_name,

            request.digest_function = ?grpc_request.get_ref().digest_function

        )

    )]

    async fn find_missing_blobs(

        &self,

        grpc_request: Request<FindMissingBlobsRequest>,

    ) -> Result<Response<FindMissingBlobsResponse>, Status> {

        let request = grpc_request.into_inner();

        let digest_function = request.digest_function;

        self.inner_find_missing_blobs(request)

            .instrument(error_span!("cas_server_find_missing_blobs"))

            .with_context(

                make_ctx_for_hash_func(digest_function)

                    .err_tip(|| "In CasServer::find_missing_blobs")?,

            )

            .await

            .err_tip(|| "Failed on find_missing_blobs() command")

            .map_err(Into::into)

    }

    #[instrument(

        err,

        ret(level = Level::DEBUG),

        level = Level::ERROR,

        skip_all,

        fields(request = ?grpc_request.get_ref())

    )]

    async fn batch_update_blobs(

        &self,

        grpc_request: Request<BatchUpdateBlobsRequest>,

    ) -> Result<Response<BatchUpdateBlobsResponse>, Status> {

        let request = grpc_request.into_inner();

        let digest_function = request.digest_function;

        self.inner_batch_update_blobs(request)

            .instrument(error_span!("cas_server_batch_update_blobs"))

            .with_context(

                make_ctx_for_hash_func(digest_function)

                    .err_tip(|| "In CasServer::batch_update_blobs")?,

            )

            .await

            .err_tip(|| "Failed on batch_update_blobs() command")

            .map_err(Into::into)

    }

    #[instrument(

        err,

        ret(level = Level::INFO),

        level = Level::ERROR,

        skip_all,

        fields(request = ?grpc_request.get_ref())

    )]

    async fn batch_read_blobs(

        &self,

        grpc_request: Request<BatchReadBlobsRequest>,

    ) -> Result<Response<BatchReadBlobsResponse>, Status> {

        let request = grpc_request.into_inner();

        let digest_function = request.digest_function;

        self.inner_batch_read_blobs(request)

            .instrument(error_span!("cas_server_batch_read_blobs"))

            .with_context(

                make_ctx_for_hash_func(digest_function)

                    .err_tip(|| "In CasServer::batch_read_blobs")?,

            )

            .await

            .err_tip(|| "Failed on batch_read_blobs() command")

            .map_err(Into::into)

    }

    #[instrument(

        err,

        level = Level::ERROR,

        skip_all,

        fields(request = ?grpc_request.get_ref())

    )]

    async fn get_tree(

        &self,

        grpc_request: Request<GetTreeRequest>,

    ) -> Result<Response<Self::GetTreeStream>, Status> {

        let request = grpc_request.into_inner();

        let digest_function = request.digest_function;

        let resp = self

            .inner_get_tree(request)

            .instrument(error_span!("cas_server_get_tree"))

            .with_context(

                make_ctx_for_hash_func(digest_function).err_tip(|| "In CasServer::get_tree")?,

            )

            .await

            .err_tip(|| "Failed on get_tree() command")

            .map(|stream| -> Response<Self::GetTreeStream> { Response::new(Box::pin(stream)) })

            .map_err(Into::into);

        if resp.is_ok() {

            debug!(return = "Ok(<stream>)");

        }

        resp

    }

}