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

//

// Licensed under the Apache License, Version 2.0 (the "License");

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

// You may obtain a copy of the License at

//

//    http://www.apache.org/licenses/LICENSE-2.0

//

// 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::cmp;

use std::pin::Pin;

use std::sync::Arc;

use async_trait::async_trait;

use bincode::config::{FixintEncoding, WithOtherIntEncoding};

use bincode::{DefaultOptions, Options};

use futures::stream::{self, FuturesOrdered, StreamExt, TryStreamExt};

use nativelink_config::stores::DedupSpec;

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

use nativelink_metric::MetricsComponent;

use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};

use nativelink_util::common::DigestInfo;

use nativelink_util::fastcdc::FastCDC;

use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator};

use nativelink_util::store_trait::{Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo};

use serde::{Deserialize, Serialize};

use tokio_util::codec::FramedRead;

use tokio_util::io::StreamReader;

use tracing::{Level, event};

use crate::cas_utils::is_zero_digest;

// NOTE: If these change update the comments in `stores.rs` to reflect

// the new defaults.

const DEFAULT_MIN_SIZE: u64 = 64 * 1024;

const DEFAULT_NORM_SIZE: u64 = 256 * 1024;

const DEFAULT_MAX_SIZE: u64 = 512 * 1024;

const DEFAULT_MAX_CONCURRENT_FETCH_PER_GET: usize = 10;

#[derive(Serialize, Deserialize, PartialEq, Debug, Default, Clone)]

pub struct DedupIndex {

    pub entries: Vec<DigestInfo>,

}

#[derive(MetricsComponent)]

pub struct DedupStore {

    #[metric(group = "index_store")]

    index_store: Store,

    #[metric(group = "content_store")]

    content_store: Store,

    fast_cdc_decoder: FastCDC,

    #[metric(help = "Maximum number of concurrent fetches per get")]

    max_concurrent_fetch_per_get: usize,

    bincode_options: WithOtherIntEncoding<DefaultOptions, FixintEncoding>,

}

impl std::fmt::Debug for DedupStore {

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        f.debug_struct("DedupStore")

            .field("index_store", &self.index_store)

            .field("content_store", &self.content_store)

            .field("fast_cdc_decoder", &self.fast_cdc_decoder)

            .field(

                "max_concurrent_fetch_per_get",

                &self.max_concurrent_fetch_per_get,

            )

            .finish_non_exhaustive()

    }

}

impl DedupStore {

    pub fn new(

        spec: &DedupSpec,

        index_store: Store,

        content_store: Store,

    ) -> Result<Arc<Self>, Error> {

        let min_size = if spec.min_size == 0 {

  Branch (82:27): [True: 0, False: 8]

  Branch (82:27): [Folded - Ignored]

            DEFAULT_MIN_SIZE

        } else {

            u64::from(spec.min_size)

        };

        let normal_size = if spec.normal_size == 0 {

  Branch (87:30): [True: 0, False: 8]

  Branch (87:30): [Folded - Ignored]

            DEFAULT_NORM_SIZE

        } else {

            u64::from(spec.normal_size)

        };

        let max_size = if spec.max_size == 0 {

  Branch (92:27): [True: 0, False: 8]

  Branch (92:27): [Folded - Ignored]

            DEFAULT_MAX_SIZE

        } else {

            u64::from(spec.max_size)

        };

        let max_concurrent_fetch_per_get = if spec.max_concurrent_fetch_per_get == 0 {

  Branch (97:47): [True: 0, False: 8]

  Branch (97:47): [Folded - Ignored]

            DEFAULT_MAX_CONCURRENT_FETCH_PER_GET

        } else {

            spec.max_concurrent_fetch_per_get as usize

        };

        Ok(Arc::new(Self {

            index_store,

            content_store,

            fast_cdc_decoder: FastCDC::new(

                usize::try_from(min_size).err_tip(|| 
"Could not convert min_size to usize"0
)
?0
,

                usize::try_from(normal_size)

                    .err_tip(|| 
"Could not convert normal_size to usize"0
)
?0
,

                usize::try_from(max_size).err_tip(|| 
"Could not convert max_size to usize"0
)
?0
,

            ),

            max_concurrent_fetch_per_get,

            bincode_options: DefaultOptions::new().with_fixint_encoding(),

        }))

    }

    async fn has(self: Pin<&Self>, key: StoreKey<'_>) -> Result<Option<u64>, Error> {

        // First we need to load the index that contains where the individual parts actually

        // can be fetched from.

        let index_entries = {

            let maybe_data = self

                .index_store

                .get_part_unchunked(key.borrow(), 0, None)

                .await

                .err_tip(|| 
"Failed to read index store in dedup store"1
);

            let 
data3
 = match maybe_data {

                Err(e) => {

                    if e.code == Code::NotFound {

  Branch (127:24): [True: 1, False: 0]

  Branch (127:24): [Folded - Ignored]

                        return Ok(None);

                    }

                    return Err(e);

                }

                Ok(data) => data,

            };

            match self.bincode_options.deserialize::<DedupIndex>(&data) {

                Err(err) => {

                    event!(

                        Level::WARN,

                        ?key,

                        ?err,

                        "Failed to deserialize index in dedup store",

                    );

                    // We return the equivalent of NotFound here so the client is happy.

                    return Ok(None);

                }

                Ok(v) => v,

            }

        };

        let digests: Vec<_> = index_entries

            .entries

            .into_iter()

            .map(StoreKey::Digest)

            .collect();

        let mut sum = 0;

        for size in 
self.content_store.has_many(&digests).await3
?0
 {

            let Some(
size7
) = size else {

  Branch (157:17): [True: 7, False: 1]

  Branch (157:17): [Folded - Ignored]

                // A part is missing so return None meaning not-found.

                // This will abort all in-flight queries related to this request.

                return Ok(None);

            };

            sum += size;

        }

        Ok(Some(sum))

    }

}

#[async_trait]

impl StoreDriver for DedupStore {

    async fn has_with_results(

        self: Pin<&Self>,

        digests: &[StoreKey<'_>],

        results: &mut [Option<u64>],

    ) -> Result<(), Error> {

        digests

            .iter()

            .zip(results.iter_mut())

            .map(|(key, result)| async move {

                if is_zero_digest(key.borrow()) {

  Branch (179:20): [True: 1, False: 4]

  Branch (179:20): [Folded - Ignored]

                    *result = Some(0);

                    return Ok(());

                }

                match self.has(key.borrow()).await {

                    Ok(maybe_size) => {

                        *result = maybe_size;

                        Ok(())

                    }

                    Err(err) => Err(err),

                }

            })

            .collect::<FuturesOrdered<_>>()

            .try_collect()

            .await

    }

    async fn update(

        self: Pin<&Self>,

        key: StoreKey<'_>,

        reader: DropCloserReadHalf,

        _size_info: UploadSizeInfo,

    ) -> Result<(), Error> {

        let mut bytes_reader = StreamReader::new(reader);

        let frame_reader = FramedRead::new(&mut bytes_reader, self.fast_cdc_decoder.clone());

        let index_entries = frame_reader

            .map(|r| r.err_tip(|| 
"Failed to decode frame from fast_cdc"0
))

            .
map_ok7
(|frame| async move {

                let hash = blake3::hash(&frame[..]).into();

                let index_entry = DigestInfo::new(hash, frame.len() as u64);

                if self

  Branch (210:20): [True: 0, False: 83]

  Branch (210:20): [Folded - Ignored]

                    .content_store

                    .has(index_entry)

                    .await

                    .err_tip(|| 
"Failed to call .has() in DedupStore::update()"0
)
?0

                    .is_some()

                {

                    // If our store has this digest, we don't need to upload it.

                    return Result::<_, Error>::Ok(index_entry);

                }

                self.content_store

                    .update_oneshot(index_entry, frame)

                    .await

                    .err_tip(|| 
"Failed to update content store in dedup_store"0
)
?0
;

                Ok(index_entry)

            })

            .try_buffered(self.max_concurrent_fetch_per_get)

            .try_collect()

            .await
?0
;

        let serialized_index = self

            .bincode_options

            .serialize(&DedupIndex {

                entries: index_entries,

            })

            .map_err(|e| {

                make_err!(

                    Code::Internal,

                    "Failed to serialize index in dedup_store : {:?}",

                    e

                )

            })?;

        self.index_store

            .update_oneshot(key, serialized_index.into())

            .await

            .err_tip(|| 
"Failed to insert our index entry to index_store in dedup_store"0
)
?0
;

        Ok(())

    }

    async fn get_part(

        self: Pin<&Self>,

        key: StoreKey<'_>,

        writer: &mut DropCloserWriteHalf,

        offset: u64,

        length: Option<u64>,

    ) -> Result<(), Error> {

        // Special case for if a client tries to read zero bytes.

        if length == Some(0) {

  Branch (259:12): [True: 30, False: 905]

  Branch (259:12): [Folded - Ignored]

            writer

                .send_eof()

                .err_tip(|| 
"Failed to write EOF out from get_part dedup"0
)
?0
;

            return Ok(());

        }

        // First we need to download the index that contains where the individual parts actually

        // can be fetched from.

        let index_entries = {

            let data = self

                .index_store

                .get_part_unchunked(key, 0, None)

                .await

                .err_tip(|| 
"Failed to read index store in dedup store"0
)
?0
;

            self.bincode_options

                .deserialize::<DedupIndex>(&data)

                .map_err(|e| {

                    make_err!(

                        Code::Internal,

                        "Failed to deserialize index in dedup_store::get_part : {:?}",

                        e

                    )

                })?

        };

        let mut start_byte_in_stream: u64 = 0;

        let entries = {

            if offset == 0 && 
length.is_none31
() {

  Branch (287:16): [True: 31, False: 874]
  Branch (287:31): [True: 2, False: 29]

  Branch (287:16): [Folded - Ignored]
  Branch (287:31): [Folded - Ignored]

                index_entries.entries

            } else {

                let mut current_entries_sum = 0;

                let mut entries = Vec::with_capacity(index_entries.entries.len());

                for 
entry4.24k
 in index_entries.entries {

                    let first_byte = current_entries_sum;

                    let entry_size = entry.size_bytes();

                    current_entries_sum += entry_size;

                    // Filter any items who's end byte is before the first requested byte.

                    if current_entries_sum <= offset {

  Branch (297:24): [True: 1.86k, False: 2.38k]

  Branch (297:24): [Folded - Ignored]

                        start_byte_in_stream = current_entries_sum;

                        continue;

                    }

                    // If we are not going to read any bytes past the length we are done.

                    if let Some(
length2.37k
) = length {

  Branch (302:28): [True: 2.37k, False: 8]

  Branch (302:28): [Folded - Ignored]

                        if first_byte >= offset + length {

  Branch (303:28): [True: 326, False: 2.04k]

  Branch (303:28): [Folded - Ignored]

                            break;

                        }

                    }

                    entries.push(entry);

                }

                entries

            }

        };

        // Second we we create a stream of futures for each chunk, but buffer/limit them so only

        // `max_concurrent_fetch_per_get` will be executed at a time.

        // The results will be streamed out in the same order they are in the entries table.

        // The results will execute in a "window-like" fashion, meaning that if we limit to

        // 5 requests at a time, and request 3 is stalled, request 1 & 2 can be output and

        // request 4 & 5 can be executing (or finished) while waiting for 3 to finish.

        // Note: We will buffer our data here up to:

        // `spec.max_size * spec.max_concurrent_fetch_per_get` per `get_part()` request.

        let mut entries_stream = stream::iter(entries)

            .map(move |index_entry| async move {

                let 
data2.11k
 = self

                    .content_store

                    .get_part_unchunked(index_entry, 0, None)

                    .await

                    .err_tip(|| 
"Failed to get_part in content_store in dedup_store"1
)
?1
;

                Result::<_, Error>::Ok(data)

            })

            .buffered(self.max_concurrent_fetch_per_get);

        // Stream out the buffered data one at a time and write the data to our writer stream.

        // In the event any of these error, we will abort early and abandon all the rest of the

        // streamed data.

        // Note: Need to take special care to ensure we send the proper slice of data requested.

        let mut bytes_to_skip = usize::try_from(offset - start_byte_in_stream)

            .err_tip(|| 
"Could not convert (offset - start_byte_in_stream) to usize"0
)
?0
;

        let mut bytes_to_send = usize::try_from(length.unwrap_or(u64::MAX - offset))

            .err_tip(|| 
"Could not convert length to usize"0
)
?0
;

        while let Some(
result2.11k
) = entries_stream.next().await {

  Branch (341:19): [True: 2.11k, False: 904]

  Branch (341:19): [Folded - Ignored]

            let 
mut data2.11k
 = result.err_tip(|| 
"Inner store iterator closed early in DedupStore"1
)
?1
;

            assert!(

                bytes_to_skip <= data.len(),

                "Formula above must be wrong, {} > {}",

                bytes_to_skip,

                data.len()

            );

            let end_pos = cmp::min(data.len(), bytes_to_send + bytes_to_skip);

            if bytes_to_skip != 0 || 
data.len() > bytes_to_send1.38k
 {

  Branch (350:16): [True: 728, False: 1.38k]
  Branch (350:38): [True: 324, False: 1.06k]

  Branch (350:16): [Folded - Ignored]
  Branch (350:38): [Folded - Ignored]

                data = data.slice(bytes_to_skip..end_pos);

            }

            writer

                .send(data)

                .await

                .err_tip(|| 
"Failed to write data to get_part dedup"0
)
?0
;

            bytes_to_send -= end_pos - bytes_to_skip;

            bytes_to_skip = 0;

        }

        // Finish our stream by writing our EOF and shutdown the stream.

        writer

            .send_eof()

            .err_tip(|| 
"Failed to write EOF out from get_part dedup"0
)
?0
;

        Ok(())

    }

    fn inner_store(&self, _digest: Option<StoreKey>) -> &dyn StoreDriver {

        self

    }

    fn as_any<'a>(&'a self) -> &'a (dyn std::any::Any + Sync + Send + 'static) {

        self

    }

    fn as_any_arc(self: Arc<Self>) -> Arc<dyn std::any::Any + Sync + Send + 'static> {

        self

    }

}

default_health_status_indicator!(DedupStore);