// 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 std::collections::HashMap;

use std::sync::Arc;

use nativelink_config::cas_server::{

    CapabilitiesConfig, CasStoreConfig, InstanceName, WithInstanceName,

};

use nativelink_error::{Error, ResultExt};

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

    Capabilities, CapabilitiesServer as Server,

};

use nativelink_proto::build::bazel::remote::execution::v2::digest_function::Value as DigestFunction;

use nativelink_proto::build::bazel::remote::execution::v2::priority_capabilities::PriorityRange;

use nativelink_proto::build::bazel::remote::execution::v2::symlink_absolute_path_strategy::Value as SymlinkAbsolutePathStrategy;

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

    ActionCacheUpdateCapabilities, CacheCapabilities, ExecutionCapabilities, FastCdc2020Params,

    GetCapabilitiesRequest, PriorityCapabilities, ServerCapabilities, compressor,

};

use nativelink_proto::build::bazel::semver::SemVer;

use nativelink_scheduler::known_platform_property_provider::KnownPlatformPropertyProvider;

use nativelink_util::digest_hasher::default_digest_hasher_func;

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

use tracing::{Level, instrument};

use crate::wire_compression::RemoteCacheCompressionInstances;

const MAX_BATCH_TOTAL_SIZE: i64 = 64 * 1024;

#[derive(Debug, Default)]

pub struct CapabilitiesServer {

    supported_node_properties_for_instance: HashMap<InstanceName, Vec<String>>,

    // Kept separate from `supported_node_properties_for_instance`: that map is sent

    // verbatim to clients as `ExecutionCapabilities.supported_node_properties` (REAPI

    // file metadata like mtime/unix mode), while compression is advertised through

    // `CacheCapabilities.supported_compressors`. Merging them would leak a fake node

    // property onto the wire.

    remote_cache_compression_instances: RemoteCacheCompressionInstances,

    chunking_params_for_instance: HashMap<InstanceName, FastCdc2020Params>,

}

impl CapabilitiesServer {

    pub async fn new(

        configs: &[WithInstanceName<CapabilitiesConfig>],

        scheduler_map: &HashMap<String, Arc<dyn KnownPlatformPropertyProvider>>,

        remote_cache_compression_instances: &RemoteCacheCompressionInstances,

        cas_configs: &[WithInstanceName<CasStoreConfig>],

    ) -> Result<Self, Error> {

        let mut chunking_params_for_instance = HashMap::new();

        for 
cas_config0
 in cas_configs {

            if let Some(chunking_config) = &cas_config.experimental_chunking {

                let avg_chunk_size_bytes = chunking_config

                    .validated_avg_chunk_size_bytes()

                    .err_tip(|| {

                        format!(

                            "In 'experimental_chunking' of instance '{}'",

                            cas_config.instance_name

                        )

                    })?;

                chunking_params_for_instance.insert(

                    cas_config.instance_name.clone(),

                    FastCdc2020Params {

                        avg_chunk_size_bytes,

                        seed: 0,

                    },

                );

            }

        }

        let mut supported_node_properties_for_instance = HashMap::new();

        for config in configs {

            if let Some(
remote_execution_cfg1
) = &config.remote_execution {

                let scheduler =

                    scheduler_map

                        .get(&remote_execution_cfg.scheduler)

                        .err_tip(|| 
{0

                            format!(

                                "Scheduler needs config for '{}' because it exists in capabilities",

                                remote_execution_cfg.scheduler

                            )

                        })?;

                let properties = scheduler

                    .get_known_properties(&config.instance_name)

                    .await

                    .err_tip(|| 
{0

                        format!(

                            "Failed to get platform properties for {}",

                            config.instance_name

                        )

                    })?;

                supported_node_properties_for_instance

                    .insert(config.instance_name.clone(), properties);

            }

        }

        Ok(Self {

            supported_node_properties_for_instance,

            remote_cache_compression_instances: remote_cache_compression_instances.clone(),

            chunking_params_for_instance,

        })

    }

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

        Server::new(self)

    }

}

#[tonic::async_trait]

impl Capabilities for CapabilitiesServer {

    #[instrument(

        err,

        ret(level = Level::INFO),

        level = Level::ERROR,

        skip_all,

        fields(request = ?grpc_request.get_ref())

    )]

    async fn get_capabilities(

        &self,

        grpc_request: Request<GetCapabilitiesRequest>,

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

        let request = grpc_request.into_inner();

        let instance_name = request.instance_name;

        let maybe_supported_node_properties = self

            .supported_node_properties_for_instance

            .get(&instance_name);

        let execution_capabilities =

            maybe_supported_node_properties.map(|props_for_instance| ExecutionCapabilities {

                digest_function: default_digest_hasher_func().proto_digest_func().into(),

                exec_enabled: true, // TODO(palfrey) Make this configurable.

                execution_priority_capabilities: Some(PriorityCapabilities {

                    priorities: vec![PriorityRange {

                        min_priority: 0,

                        max_priority: i32::MAX,

                    }],

                }),

                supported_node_properties: props_for_instance.clone(),

                digest_functions: vec![

                    DigestFunction::Sha256.into(),

                    DigestFunction::Blake3.into(),

                ],

            });

        let supported_compressors = if self

            .remote_cache_compression_instances

            .enabled_for(&instance_name)

        {

            vec![compressor::Value::Zstd.into()]

        } else {

            Vec::new()

        };

        let chunking_params = self.chunking_params_for_instance.get(&instance_name);

        let resp = ServerCapabilities {

            cache_capabilities: Some(CacheCapabilities {

                digest_functions: vec![

                    DigestFunction::Sha256.into(),

                    DigestFunction::Blake3.into(),

                ],

                action_cache_update_capabilities: Some(ActionCacheUpdateCapabilities {

                    update_enabled: true,

                }),

                cache_priority_capabilities: None,

                max_batch_total_size_bytes: MAX_BATCH_TOTAL_SIZE,

                symlink_absolute_path_strategy: SymlinkAbsolutePathStrategy::Disallowed.into(),

                supported_compressors: supported_compressors.clone(),

                supported_batch_update_compressors: supported_compressors,

                max_cas_blob_size_bytes: 0,

                split_blob_support: chunking_params.is_some(),

                splice_blob_support: chunking_params.is_some(),

                fast_cdc_2020_params: chunking_params.copied(),

                rep_max_cdc_params: None,

            }),

            execution_capabilities,

            deprecated_api_version: None,

            low_api_version: Some(SemVer {

                major: 2,

                minor: 0,

                patch: 0,

                prerelease: String::new(),

            }),

            high_api_version: Some(SemVer {

                major: 2,

                minor: 3,

                patch: 0,

                prerelease: String::new(),

            }),

        };

        Ok(Response::new(resp))

    }

}