qdrant-edge 0.7.0

use std::path::{Path, PathBuf};

use crate::common::counter::hardware_counter::HardwareCounterCell;
use crate::common::generic_consts::{Random, Sequential};
use crate::common::types::PointOffsetType;
use fs_err as fs;
use crate::gridstore::config::StorageOptions;
use crate::gridstore::{Blob, Gridstore};
use serde_json::Value;

use crate::segment::common::Flusher;
use crate::segment::common::operation_error::{OperationError, OperationResult};
use crate::segment::json_path::JsonPath;
use crate::segment::payload_storage::{PayloadStorage, PayloadStorageRead};
use crate::segment::types::{OwnedPayloadRef, Payload, PayloadKeyTypeRef};

const STORAGE_PATH: &str = "payload_storage";

impl Blob for Payload {
    fn to_bytes(&self) -> Vec<u8> {
        serde_json::to_vec(self).unwrap()
    }

    fn from_bytes(data: &[u8]) -> Self {
        serde_json::from_slice(data).unwrap()
    }
}

#[derive(Debug)]
pub struct MmapPayloadStorage {
    storage: Gridstore<Payload>,
    populate: bool,
}

impl MmapPayloadStorage {
    pub fn open_or_create(path: PathBuf, populate: bool) -> OperationResult<Self> {
        let path = storage_dir(path);
        if path.exists() {
            Self::open(path, populate)
        } else {
            // create folder if it does not exist
            fs::create_dir_all(&path).map_err(|_| {
                OperationError::service_error("Failed to create mmap payload storage directory")
            })?;
            Ok(Self::new(path, populate)?)
        }
    }

    fn open(path: PathBuf, populate: bool) -> OperationResult<Self> {
        let storage = Gridstore::open(path).map_err(|err| {
            OperationError::service_error(format!("Failed to open mmap payload storage: {err}"))
        })?;

        if populate {
            storage.populate()?;
        }

        Ok(Self { storage, populate })
    }

    fn new(path: PathBuf, populate: bool) -> OperationResult<Self> {
        let storage = Gridstore::new(path, StorageOptions::default())?;

        if populate {
            storage.populate()?;
        }

        Ok(Self { storage, populate })
    }

    /// Populate all pages in the mmap.
    /// Block until all pages are populated.
    pub fn populate(&self) -> OperationResult<()> {
        self.storage.populate()?;
        Ok(())
    }

    /// Drop disk cache.
    pub fn clear_cache(&self) -> OperationResult<()> {
        self.storage.clear_cache()?;
        Ok(())
    }
}

impl PayloadStorageRead for MmapPayloadStorage {
    fn get(
        &self,
        point_offset: PointOffsetType,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<Payload> {
        match self.storage.get_value::<Random>(point_offset, hw_counter)? {
            Some(payload) => Ok(payload),
            None => Ok(Default::default()),
        }
    }

    fn get_sequential(
        &self,
        point_offset: PointOffsetType,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<Payload> {
        match self
            .storage
            .get_value::<Sequential>(point_offset, hw_counter)?
        {
            Some(payload) => Ok(payload),
            None => Ok(Default::default()),
        }
    }

    fn payload_ref(
        &self,
        point_offset: PointOffsetType,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<OwnedPayloadRef<'_>> {
        let payload = self.get(point_offset, hw_counter)?;
        Ok(OwnedPayloadRef::from(payload))
    }

    fn iter<F>(&self, mut callback: F, hw_counter: &HardwareCounterCell) -> OperationResult<()>
    where
        F: FnMut(PointOffsetType, &Payload) -> OperationResult<bool>,
    {
        self.storage.iter(
            |point_id, payload| callback(point_id, &payload),
            hw_counter.ref_payload_io_read_counter(),
        )
    }

    fn get_storage_size_bytes(&self) -> OperationResult<usize> {
        Ok(self.storage.get_storage_size_bytes()?)
    }

    fn is_on_disk(&self) -> bool {
        !self.populate
    }
}

impl PayloadStorage for MmapPayloadStorage {
    fn overwrite(
        &mut self,
        point_id: PointOffsetType,
        payload: &Payload,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<()> {
        self.storage
            .put_value(point_id, payload, hw_counter.ref_payload_io_write_counter())?;
        Ok(())
    }

    fn set(
        &mut self,
        point_id: PointOffsetType,
        payload: &Payload,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<()> {
        match self.storage.get_value::<Random>(point_id, hw_counter)? {
            Some(mut point_payload) => {
                point_payload.merge(payload);
                self.storage.put_value(
                    point_id,
                    &point_payload,
                    hw_counter.ref_payload_io_write_counter(),
                )?;
            }
            None => {
                self.storage.put_value(
                    point_id,
                    payload,
                    hw_counter.ref_payload_io_write_counter(),
                )?;
            }
        }
        Ok(())
    }

    fn set_by_key(
        &mut self,
        point_id: PointOffsetType,
        payload: &Payload,
        key: &JsonPath,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<()> {
        match self.storage.get_value::<Random>(point_id, hw_counter)? {
            Some(mut point_payload) => {
                point_payload.merge_by_key(payload, key);
                self.storage.put_value(
                    point_id,
                    &point_payload,
                    hw_counter.ref_payload_io_write_counter(),
                )?;
            }
            None => {
                let mut dest_payload = Payload::default();
                dest_payload.merge_by_key(payload, key);
                self.storage.put_value(
                    point_id,
                    &dest_payload,
                    hw_counter.ref_payload_io_write_counter(),
                )?;
            }
        }
        Ok(())
    }

    fn delete(
        &mut self,
        point_id: PointOffsetType,
        key: PayloadKeyTypeRef,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<Vec<Value>> {
        match self.storage.get_value::<Random>(point_id, hw_counter)? {
            Some(mut payload) => {
                let res = payload.remove(key);
                if !res.is_empty() {
                    self.storage.put_value(
                        point_id,
                        &payload,
                        hw_counter.ref_payload_io_write_counter(),
                    )?;
                }
                Ok(res)
            }
            None => Ok(vec![]),
        }
    }

    fn clear(
        &mut self,
        point_id: PointOffsetType,
        _: &HardwareCounterCell,
    ) -> OperationResult<Option<Payload>> {
        let res = self.storage.delete_value(point_id)?;
        Ok(res)
    }

    #[cfg(test)]
    fn clear_all(&mut self, _: &HardwareCounterCell) -> OperationResult<()> {
        self.storage.clear().map_err(|err| {
            OperationError::service_error(format!("Failed to clear mmap payload storage: {err}"))
        })
    }

    fn flusher(&self) -> Flusher {
        let storage_flusher = self.storage.flusher();
        Box::new(move || {
            storage_flusher().map_err(|err| {
                OperationError::service_error(format!(
                    "Failed to flush mmap payload gridstore: {err}"
                ))
            })
        })
    }

    fn files(&self) -> Vec<PathBuf> {
        self.storage.files()
    }

    fn immutable_files(&self) -> Vec<PathBuf> {
        self.storage.immutable_files()
    }
}

/// Get storage directory for this payload storage
pub fn storage_dir<P: AsRef<Path>>(segment_path: P) -> PathBuf {
    segment_path.as_ref().join(STORAGE_PATH)
}