// Copyright 2024 KVCache.AI
//
// 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.

//! Safe Rust wrapper around the Mooncake Store C API (`store_c.h`).

use std::ffi::{c_void, CString};

use crate::error::StoreError;

// Raw FFI bindings generated by build.rs / bindgen.
#[allow(dead_code)]
mod ffi {
    #![allow(non_snake_case)]
    #![allow(non_camel_case_types)]
    #![allow(non_upper_case_globals)]
    #![allow(clippy::useless_transmute)]
    include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
}

// ---------------------------------------------------------------------------
// ReplicateConfig
// ---------------------------------------------------------------------------

/// Replication settings for a put operation.
///
/// This is the Rust counterpart of `mooncake_replicate_config_t` from
/// `store_c.h`.
#[derive(Debug, Clone, Default)]
pub struct ReplicateConfig {
    /// Number of replicas to create (0 means use the server default).
    pub replica_num: usize,
    /// Prefer a replica on the same NUMA node / host ("soft pin").
    pub with_soft_pin: bool,
    /// Declare whether the object will never be evicted.
    pub with_hard_pin: bool,
    /// Whitelist of segment names that should host a replica.
    pub preferred_segments: Vec<String>,
}

impl ReplicateConfig {
    /// Converts this config into the C representation together with the
    /// owning `CString` objects that back the segment-name pointers.
    ///
    /// The returned `_strings` must stay alive for as long as
    /// `ffi::mooncake_replicate_config_t` is in use.
    fn to_ffi(
        &self,
    ) -> Result<
        (ffi::mooncake_replicate_config_t, Vec<CString>, Vec<*const libc::c_char>),
        StoreError,
    > {
        let strings: Vec<CString> = self
            .preferred_segments
            .iter()
            .map(|s| CString::new(s.as_str()).map_err(StoreError::from))
            .collect::<Result<_, _>>()?;

        let ptrs: Vec<*const libc::c_char> = strings.iter().map(|s| s.as_ptr()).collect();

        let c_config = ffi::mooncake_replicate_config_t {
            replica_num: self.replica_num,
            with_soft_pin: i32::from(self.with_soft_pin),
            with_hard_pin: i32::from(self.with_hard_pin),
            preferred_segments: if ptrs.is_empty() {
                std::ptr::null_mut()
            } else {
                ptrs.as_ptr() as *mut *const libc::c_char
            },
            preferred_segments_count: ptrs.len(),
        };

        Ok((c_config, strings, ptrs))
    }
}

// ---------------------------------------------------------------------------
// MooncakeStore
// ---------------------------------------------------------------------------

/// A handle to a Mooncake distributed KV-cache store.
///
/// # Lifecycle
///
/// ```ignore
/// let store = MooncakeStore::new()?;
/// store.setup("localhost", "http://127.0.0.1:8080/metadata",
///             512 << 20, 128 << 20, "tcp", "", "127.0.0.1:50051")?;
/// store.put("my-key", b"hello world", None)?;
/// let mut buf = vec![0u8; 11];
/// store.get_into("my-key", buf.as_mut_ptr() as *mut c_void, buf.len())?;
/// ```
///
/// The store is automatically destroyed when dropped.
pub struct MooncakeStore {
    handle: ffi::mooncake_store_t,
}

// The underlying C object is internally synchronised; sharing the handle
// across threads is safe.
unsafe impl Send for MooncakeStore {}
unsafe impl Sync for MooncakeStore {}

impl MooncakeStore {
    // -----------------------------------------------------------------------
    // Lifecycle
    // -----------------------------------------------------------------------

    /// Allocate a new (uninitialised) store handle.
    ///
    /// Call [`MooncakeStore::setup`] before performing any data operations.
    pub fn new() -> Result<Self, StoreError> {
        let handle = unsafe { ffi::mooncake_store_create() };
        if handle.is_null() {
            return Err(StoreError::NullHandle);
        }
        Ok(Self { handle })
    }

    /// Initialise the store client.
    ///
    /// # Parameters
    /// - `local_hostname` – IP or hostname of *this* node.
    /// - `metadata_server` – URL of the metadata server
    ///   (e.g. `"http://127.0.0.1:8080/metadata"` or `"etcd://127.0.0.1:2379"`).
    /// - `global_segment_size` – size of each RDMA memory segment in bytes.
    /// - `local_buffer_size` – size of the local transfer buffer in bytes.
    /// - `protocol` – transport protocol (`"tcp"`, `"rdma"`, …).
    /// - `device_name` – network device name (empty string = auto-select).
    /// - `master_server_addr` – address of the Mooncake master service
    ///   (e.g. `"127.0.0.1:50051"`).
    pub fn setup(
        &self,
        local_hostname: &str,
        metadata_server: &str,
        global_segment_size: u64,
        local_buffer_size: u64,
        protocol: &str,
        device_name: &str,
        master_server_addr: &str,
    ) -> Result<(), StoreError> {
        let local_hostname_c = CString::new(local_hostname)?;
        let metadata_server_c = CString::new(metadata_server)?;
        let protocol_c = CString::new(protocol)?;
        let device_name_c = CString::new(device_name)?;
        let master_server_addr_c = CString::new(master_server_addr)?;

        let rc = unsafe {
            ffi::mooncake_store_setup(
                self.handle,
                local_hostname_c.as_ptr(),
                metadata_server_c.as_ptr(),
                global_segment_size,
                local_buffer_size,
                protocol_c.as_ptr(),
                device_name_c.as_ptr(),
                master_server_addr_c.as_ptr(),
            )
        };
        if rc != 0 {
            return Err(StoreError::OperationFailed(rc));
        }
        Ok(())
    }

    /// Perform a connectivity health-check.
    ///
    /// Returns `Ok(())` when the store is reachable, or an error otherwise.
    pub fn health_check(&self) -> Result<(), StoreError> {
        let rc = unsafe { ffi::mooncake_store_health_check(self.handle) };
        if rc != 0 {
            return Err(StoreError::OperationFailed(rc));
        }
        Ok(())
    }

    // -----------------------------------------------------------------------
    // Put operations
    // -----------------------------------------------------------------------

    /// Store `value` under `key`.
    ///
    /// The data is copied from `value` into the store's internal buffers, so
    /// `value` does not need to be pinned or registered.
    pub fn put(
        &self,
        key: &str,
        value: &[u8],
        config: Option<&ReplicateConfig>,
    ) -> Result<(), StoreError> {
        let key_c = CString::new(key)?;
        let (_c_config, _strings, _ptrs) = Self::prepare_config(config)?;
        let cfg_ptr = _c_config
            .as_ref()
            .map_or(std::ptr::null(), |c| c as *const _);

        let rc = unsafe {
            ffi::mooncake_store_put(
                self.handle,
                key_c.as_ptr(),
                value.as_ptr() as *const c_void,
                value.len(),
                cfg_ptr,
            )
        };
        if rc != 0 {
            return Err(StoreError::OperationFailed(rc));
        }
        Ok(())
    }

    /// Store data from a pre-registered, zero-copy buffer under `key`.
    ///
    /// # Safety
    ///
    /// `buffer` must be a pointer to at least `size` bytes of valid memory
    /// that has been registered with the store via
    /// [`MooncakeStore::register_buffer`].
    pub unsafe fn put_from(
        &self,
        key: &str,
        buffer: *mut c_void,
        size: usize,
        config: Option<&ReplicateConfig>,
    ) -> Result<(), StoreError> {
        let key_c = CString::new(key)?;
        let (_c_config, _strings, _ptrs) = Self::prepare_config(config)?;
        let cfg_ptr = _c_config
            .as_ref()
            .map_or(std::ptr::null(), |c| c as *const _);

        let rc = ffi::mooncake_store_put_from(self.handle, key_c.as_ptr(), buffer, size, cfg_ptr);
        if rc != 0 {
            return Err(StoreError::OperationFailed(rc));
        }
        Ok(())
    }

    // -----------------------------------------------------------------------
    // Get operations
    // -----------------------------------------------------------------------

    /// Retrieve the value for `key` into a pre-allocated `buffer`.
    ///
    /// Returns the number of bytes written on success.
    ///
    /// # Safety
    ///
    /// `buffer` must point to at least `size` bytes of writable, valid memory.
    pub unsafe fn get_into(
        &self,
        key: &str,
        buffer: *mut c_void,
        size: usize,
    ) -> Result<i64, StoreError> {
        let key_c = CString::new(key)?;
        let bytes = ffi::mooncake_store_get_into(self.handle, key_c.as_ptr(), buffer, size);
        if bytes < 0 {
            return Err(StoreError::OperationFailed(bytes as i32));
        }
        Ok(bytes)
    }

    /// Retrieve the value for `key` into a Rust-owned `Vec<u8>`.
    ///
    /// The buffer is allocated to the exact size of the stored value.
    pub fn get(&self, key: &str) -> Result<Vec<u8>, StoreError> {
        let size = self.get_size(key)?;
        let mut buf = vec![0u8; size as usize];
        let written =
            unsafe { self.get_into(key, buf.as_mut_ptr() as *mut c_void, buf.len())? };
        buf.truncate(written as usize);
        Ok(buf)
    }

    // -----------------------------------------------------------------------
    // Existence / size / hostname
    // -----------------------------------------------------------------------

    /// Returns `true` when `key` exists in the store.
    pub fn is_exist(&self, key: &str) -> Result<bool, StoreError> {
        let key_c = CString::new(key)?;
        let rc = unsafe { ffi::mooncake_store_is_exist(self.handle, key_c.as_ptr()) };
        match rc {
            1 => Ok(true),
            0 => Ok(false),
            _ => Err(StoreError::OperationFailed(rc)),
        }
    }

    /// Returns the size in bytes of the value stored under `key`, or an error
    /// if the key does not exist or the query fails.
    ///
    /// The underlying C API returns a single negative value (`-1`) for all
    /// error conditions (null handle, missing key, internal failure), so we
    /// report `OperationFailed` with the raw code rather than attempting to
    /// distinguish `NotFound` from other failures.
    pub fn get_size(&self, key: &str) -> Result<i64, StoreError> {
        let key_c = CString::new(key)?;
        let sz = unsafe { ffi::mooncake_store_get_size(self.handle, key_c.as_ptr()) };
        if sz < 0 {
            return Err(StoreError::OperationFailed(sz as i32));
        }
        Ok(sz)
    }

    /// Returns the hostname that the store client is registered under.
    ///
    /// Uses a 1024-byte buffer, which is well above the POSIX hostname limit
    /// of 253 characters. Returns `OperationFailed` if the hostname exceeds
    /// the buffer or the call fails for another reason.
    pub fn get_hostname(&self) -> Result<String, StoreError> {
        let mut buf = vec![0u8; 1024];
        let rc = unsafe {
            ffi::mooncake_store_get_hostname(
                self.handle,
                buf.as_mut_ptr() as *mut libc::c_char,
                buf.len(),
            )
        };
        if rc != 0 {
            return Err(StoreError::OperationFailed(rc));
        }
        let cstr = unsafe { std::ffi::CStr::from_ptr(buf.as_ptr() as *const libc::c_char) };
        Ok(cstr.to_string_lossy().into_owned())
    }

    // -----------------------------------------------------------------------
    // Remove operations
    // -----------------------------------------------------------------------

    /// Remove the entry for `key`.
    ///
    /// When `force` is `true` the key is removed even if it is currently being
    /// read by another client.
    pub fn remove(&self, key: &str, force: bool) -> Result<(), StoreError> {
        let key_c = CString::new(key)?;
        let rc = unsafe {
            ffi::mooncake_store_remove(self.handle, key_c.as_ptr(), i32::from(force))
        };
        if rc != 0 {
            return Err(StoreError::OperationFailed(rc));
        }
        Ok(())
    }

    /// Remove all keys matching `pattern` (a regular expression).
    ///
    /// Returns the number of keys removed.
    pub fn remove_by_regex(&self, pattern: &str, force: bool) -> Result<i64, StoreError> {
        let pattern_c = CString::new(pattern)?;
        let n = unsafe {
            ffi::mooncake_store_remove_by_regex(self.handle, pattern_c.as_ptr(), i32::from(force))
        };
        if n < 0 {
            return Err(StoreError::OperationFailed(n as i32));
        }
        Ok(n)
    }

    /// Remove **all** keys from the store.
    ///
    /// Returns the number of keys removed.
    pub fn remove_all(&self, force: bool) -> Result<i64, StoreError> {
        let n = unsafe { ffi::mooncake_store_remove_all(self.handle, i32::from(force)) };
        if n < 0 {
            return Err(StoreError::OperationFailed(n as i32));
        }
        Ok(n)
    }

    // -----------------------------------------------------------------------
    // Buffer registration (zero-copy operations)
    // -----------------------------------------------------------------------

    /// Register a memory region with the RDMA transport so that it can be used
    /// for zero-copy [`put_from`](MooncakeStore::put_from) /
    /// [`get_into`](MooncakeStore::get_into) operations.
    ///
    /// # Safety
    ///
    /// `buffer` must point to at least `size` bytes of valid, pinned memory
    /// that will remain valid until [`unregister_buffer`](MooncakeStore::unregister_buffer)
    /// is called.
    pub unsafe fn register_buffer(
        &self,
        buffer: *mut c_void,
        size: usize,
    ) -> Result<(), StoreError> {
        let rc = ffi::mooncake_store_register_buffer(self.handle, buffer, size);
        if rc != 0 {
            return Err(StoreError::OperationFailed(rc));
        }
        Ok(())
    }

    /// Deregister a memory region previously registered with
    /// [`register_buffer`](MooncakeStore::register_buffer).
    ///
    /// # Safety
    ///
    /// `buffer` must be the same pointer passed to `register_buffer`.
    pub unsafe fn unregister_buffer(&self, buffer: *mut c_void) -> Result<(), StoreError> {
        let rc = ffi::mooncake_store_unregister_buffer(self.handle, buffer);
        if rc != 0 {
            return Err(StoreError::OperationFailed(rc));
        }
        Ok(())
    }

    // -----------------------------------------------------------------------
    // Internal helpers
    // -----------------------------------------------------------------------

    /// Convert an `Option<&ReplicateConfig>` into a C-compatible value,
    /// returning the owning objects that keep the data alive.
    ///
    /// The caller must bind the returned tuple to a named variable in its own
    /// stack frame *before* taking a raw pointer to the
    /// `Option<mooncake_replicate_config_t>`, so that the pointer is valid for
    /// the duration of the C call.
    #[allow(clippy::type_complexity)]
    fn prepare_config(
        config: Option<&ReplicateConfig>,
    ) -> Result<
        (
            Option<ffi::mooncake_replicate_config_t>,
            Vec<CString>,
            Vec<*const libc::c_char>,
        ),
        StoreError,
    > {
        match config {
            None => Ok((None, Vec::new(), Vec::new())),
            Some(cfg) => {
                let (c_config, strings, ptrs) = cfg.to_ffi()?;
                Ok((Some(c_config), strings, ptrs))
            }
        }
    }
}

impl Drop for MooncakeStore {
    fn drop(&mut self) {
        unsafe {
            ffi::mooncake_store_destroy(self.handle);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::ffi::CStr;

    #[test]
    fn replicate_config_to_ffi_empty_segments() {
        let config = ReplicateConfig {
            replica_num: 2,
            with_soft_pin: true,
            with_hard_pin: false,
            preferred_segments: Vec::new(),
        };

        let (ffi_cfg, strings, ptrs) = config.to_ffi().expect("to_ffi should succeed");
        assert_eq!(ffi_cfg.replica_num, 2);
        assert_eq!(ffi_cfg.with_soft_pin, 1);
        assert_eq!(ffi_cfg.with_hard_pin, 0);
        assert!(ffi_cfg.preferred_segments.is_null());
        assert_eq!(ffi_cfg.preferred_segments_count, 0);
        assert!(strings.is_empty());
        assert!(ptrs.is_empty());
    }

    #[test]
    fn replicate_config_to_ffi_with_segments() {
        let config = ReplicateConfig {
            replica_num: 3,
            with_soft_pin: false,
            with_hard_pin: false,
            preferred_segments: vec!["seg-a".to_string(), "seg-b".to_string()],
        };

        let (ffi_cfg, strings, ptrs) = config.to_ffi().expect("to_ffi should succeed");
        assert_eq!(ffi_cfg.replica_num, 3);
        assert_eq!(ffi_cfg.with_soft_pin, 0);
        assert_eq!(ffi_cfg.with_hard_pin, 0);
        assert!(!ffi_cfg.preferred_segments.is_null());
        assert_eq!(ffi_cfg.preferred_segments_count, 2);
        assert_eq!(strings.len(), 2);
        assert_eq!(ptrs.len(), 2);

        let seg0 = unsafe { CStr::from_ptr(ptrs[0]) };
        let seg1 = unsafe { CStr::from_ptr(ptrs[1]) };
        assert_eq!(seg0.to_str().unwrap(), "seg-a");
        assert_eq!(seg1.to_str().unwrap(), "seg-b");
    }

    #[test]
    fn replicate_config_to_ffi_rejects_interior_nul() {
        let config = ReplicateConfig {
            replica_num: 1,
            with_soft_pin: false,
            with_hard_pin: false,
            preferred_segments: vec!["bad\0segment".to_string()],
        };

        assert!(matches!(
            config.to_ffi(),
            Err(StoreError::InvalidString(_))
        ));
    }

    #[test]
    fn prepare_config_none_returns_null_config() {
        let (c_cfg, strings, ptrs) = MooncakeStore::prepare_config(None).expect("prepare should succeed");
        assert!(c_cfg.is_none());
        assert!(strings.is_empty());
        assert!(ptrs.is_empty());
    }

    #[test]
    fn prepare_config_some_preserves_storage_and_values() {
        let config = ReplicateConfig {
            replica_num: 4,
            with_soft_pin: true,
            with_hard_pin: false,
            preferred_segments: vec!["x".to_string(), "y".to_string()],
        };

        let (c_cfg, strings, ptrs) =
            MooncakeStore::prepare_config(Some(&config)).expect("prepare should succeed");
        let cfg_ref = c_cfg.as_ref().expect("expected Some config");
        assert_eq!(cfg_ref.replica_num, 4);
        assert_eq!(cfg_ref.with_soft_pin, 1);
        assert_eq!(cfg_ref.with_hard_pin, 0);
        assert_eq!(cfg_ref.preferred_segments_count, 2);
        assert_eq!(strings.len(), 2);
        assert_eq!(ptrs.len(), 2);
    }
}
