wedb_embed 0.1.1

use crate::key_composer::slot::get_slot_id_from_key;

/// Kvrocks 1:1 内部子键结构（对标 Apache Kvrocks InternalKey）
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InternalKey<'a> {
    pub namespace: &'a [u8],
    pub slot_id: Option<u16>,
    pub key: &'a [u8],
    pub version: u64,
    pub sub_key: &'a [u8],
}

impl<'a> InternalKey<'a> {
    #[inline]
    pub fn new(
        namespace: &'a [u8],
        key: &'a [u8],
        sub_key: &'a [u8],
        version: u64,
        slot_id: Option<u16>,
    ) -> Self {
        Self {
            namespace,
            slot_id,
            key,
            version,
            sub_key,
        }
    }

    /// 从已由 ComposeNamespaceKey 编码的前缀构造 InternalKey（对标 Kvrocks InternalKey 构造函数）
    #[inline]
    pub fn from_ns_key(
        mut ns_key: &'a [u8],
        sub_key: &'a [u8],
        version: u64,
        slot_id_encoded: bool,
    ) -> Option<Self> {
        if ns_key.is_empty() {
            return None;
        }
        let ns_len = ns_key[0] as usize;
        ns_key = &ns_key[1..];
        if ns_key.len() < ns_len {
            return None;
        }
        let namespace = &ns_key[..ns_len];
        ns_key = &ns_key[ns_len..];

        let slot_id = if slot_id_encoded {
            if ns_key.len() < 2 {
                return None;
            }
            let mut s_buf = [0u8; 2];
            s_buf.copy_from_slice(&ns_key[..2]);
            ns_key = &ns_key[2..];
            Some(u16::from_be_bytes(s_buf))
        } else {
            None
        };

        let key = ns_key;
        Some(Self {
            namespace,
            slot_id,
            key,
            version,
            sub_key,
        })
    }

    /// 编码为 Kvrocks 1:1 二进制存储布局：
    /// [1B ns_len][ns][(optional 2B slot_id)][4B key_len(BE)][key][8B version(BE)][sub_key]
    #[inline]
    pub fn encode(&self) -> Vec<u8> {
        let slot_len = if self.slot_id.is_some() { 2 } else { 0 };
        let total =
            1 + self.namespace.len() + slot_len + 4 + self.key.len() + 8 + self.sub_key.len();
        let mut out = Vec::with_capacity(total);
        self.encode_to(&mut out);
        out
    }

    /// 编码到已有的缓冲区中（复用堆分配）
    #[inline]
    pub fn encode_to(&self, out: &mut Vec<u8>) {
        let slot_len = if self.slot_id.is_some() { 2 } else { 0 };
        let total =
            1 + self.namespace.len() + slot_len + 4 + self.key.len() + 8 + self.sub_key.len();
        out.reserve(total);
        out.push(self.namespace.len() as u8);
        out.extend_from_slice(self.namespace);
        if let Some(slot) = self.slot_id {
            out.extend_from_slice(&slot.to_be_bytes());
        }
        out.extend_from_slice(&(self.key.len() as u32).to_be_bytes());
        out.extend_from_slice(self.key);
        out.extend_from_slice(&self.version.to_be_bytes());
        out.extend_from_slice(self.sub_key);
    }

    /// 解码 Kvrocks 二进制内部键（零拷贝）
    #[inline]
    pub fn decode(bytes: &'a [u8], slot_id_encoded: bool) -> Option<Self> {
        if bytes.is_empty() {
            return None;
        }
        let ns_len = bytes[0] as usize;
        let mut offset = 1;
        if bytes.len() < offset + ns_len {
            return None;
        }
        let namespace = &bytes[offset..offset + ns_len];
        offset += ns_len;

        let slot_id = if slot_id_encoded {
            if bytes.len() < offset + 2 {
                return None;
            }
            let mut s_buf = [0u8; 2];
            s_buf.copy_from_slice(&bytes[offset..offset + 2]);
            offset += 2;
            Some(u16::from_be_bytes(s_buf))
        } else {
            None
        };

        if bytes.len() < offset + 4 {
            return None;
        }
        let mut klen_buf = [0u8; 4];
        klen_buf.copy_from_slice(&bytes[offset..offset + 4]);
        let key_len = u32::from_be_bytes(klen_buf) as usize;
        offset += 4;

        if bytes.len() < offset + key_len + 8 {
            return None;
        }
        let key = &bytes[offset..offset + key_len];
        offset += key_len;

        let mut ver_buf = [0u8; 8];
        ver_buf.copy_from_slice(&bytes[offset..offset + 8]);
        let version = u64::from_be_bytes(ver_buf);
        offset += 8;

        let sub_key = &bytes[offset..];

        Some(Self {
            namespace,
            slot_id,
            key,
            version,
            sub_key,
        })
    }
}

/// 构造命名空间键（对标 Kvrocks ComposeNamespaceKey）
#[inline]
pub fn compose_namespace_key(ns: &[u8], key: &[u8], slot_id: Option<u16>) -> Vec<u8> {
    let slot_len = if slot_id.is_some() { 2 } else { 0 };
    let mut out = Vec::with_capacity(1 + ns.len() + slot_len + key.len());
    out.push(ns.len() as u8);
    out.extend_from_slice(ns);
    if let Some(slot) = slot_id {
        out.extend_from_slice(&slot.to_be_bytes());
    }
    out.extend_from_slice(key);
    out
}

/// 构造命名空间键并自动根据 slot_id_encoded 计算 Slot ID（对标 Apache Kvrocks ComposeNamespaceKey）
#[inline]
pub fn compose_namespace_key_with_slot(ns: &[u8], key: &[u8], slot_id_encoded: bool) -> Vec<u8> {
    let slot_id = if slot_id_encoded {
        Some(get_slot_id_from_key(key))
    } else {
        None
    };
    compose_namespace_key(ns, key, slot_id)
}

/// 从命名空间键中零拷贝提取命名空间、槽位与用户键（对标 Kvrocks ExtractNamespaceKey）
#[inline]
pub fn extract_namespace_key(
    ns_key: &[u8],
    slot_id_encoded: bool,
) -> Option<(&[u8], Option<u16>, &[u8])> {
    if ns_key.is_empty() {
        return None;
    }
    let ns_len = ns_key[0] as usize;
    let mut offset = 1;
    if ns_key.len() < offset + ns_len {
        return None;
    }
    let ns = &ns_key[offset..offset + ns_len];
    offset += ns_len;

    let slot_id = if slot_id_encoded {
        if ns_key.len() < offset + 2 {
            return None;
        }
        let mut s_buf = [0u8; 2];
        s_buf.copy_from_slice(&ns_key[offset..offset + 2]);
        offset += 2;
        Some(u16::from_be_bytes(s_buf))
    } else {
        None
    };

    let key = &ns_key[offset..];
    Some((ns, slot_id, key))
}

/// 提取 Slot ID（对标 Kvrocks ExtractSlotId）
#[inline]
pub fn extract_slot_id(ns_key: &[u8]) -> Option<u16> {
    if ns_key.is_empty() {
        return None;
    }
    let ns_len = ns_key[0] as usize;
    let offset = 1 + ns_len;
    if ns_key.len() < offset + 2 {
        return None;
    }
    let mut s_buf = [0u8; 2];
    s_buf.copy_from_slice(&ns_key[offset..offset + 2]);
    Some(u16::from_be_bytes(s_buf))
}