wedb_embed 0.1.1

use crate::error::{Error, Result};
use crate::meta::{decode_hex_u64, u64_to_hex_16};
use crate::search::meta::{
    DistanceMetric, IndexFieldType, IndexOnDataType, VectorAlgorithm, VectorFieldMetadata,
    VectorType,
};
use sonic_rs::JsonValueTrait;
use sonic_rs::prelude::*;
use std::str;

/// 检索子键类型枚举（对标 Apache Kvrocks SearchSubkeyType）
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum SearchSubkeyType {
    IndexMeta = 0,
    Prefixes = 1,
    FieldMeta = 2,
    Field = 3,
    FieldAlias = 4,
}

impl SearchSubkeyType {
    #[inline]
    pub const fn from_u8(val: u8) -> Option<Self> {
        match val {
            0 => Some(Self::IndexMeta),
            1 => Some(Self::Prefixes),
            2 => Some(Self::FieldMeta),
            3 => Some(Self::Field),
            4 => Some(Self::FieldAlias),
            _ => None,
        }
    }
}

/// HNSW 图层级数据类型（对标 Apache Kvrocks HnswLevelType）
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum HnswLevelType {
    Node = 1,
    Edge = 2,
}

impl HnswLevelType {
    #[inline]
    pub const fn from_u8(val: u8) -> Option<Self> {
        match val {
            1 => Some(Self::Node),
            2 => Some(Self::Edge),
            _ => None,
        }
    }
}

/// 检索键构造器（对标 Apache Kvrocks SearchKey）
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SearchKey<'a> {
    pub ns: &'a str,
    pub index: &'a str,
    pub field: Option<&'a str>,
}

impl<'a> SearchKey<'a> {
    #[inline]
    pub const fn new(ns: &'a str, index: &'a str) -> Self {
        Self {
            ns,
            index,
            field: None,
        }
    }

    #[inline]
    pub const fn with_field(ns: &'a str, index: &'a str, field: &'a str) -> Self {
        Self {
            ns,
            index,
            field: Some(field),
        }
    }

    #[inline]
    pub fn put_namespace(dst: &mut Vec<u8>, ns: &str) {
        dst.push(ns.len() as u8);
        dst.extend_from_slice(ns.as_bytes());
    }

    #[inline]
    pub fn put_type(dst: &mut Vec<u8>, subkey_type: SearchSubkeyType) {
        dst.push(subkey_type as u8);
    }

    #[inline]
    pub fn put_sized_string(dst: &mut Vec<u8>, s: &str) {
        dst.extend_from_slice(&(s.len() as u32).to_be_bytes());
        dst.extend_from_slice(s.as_bytes());
    }

    #[inline]
    pub fn get_sized_string<'b>(input: &mut &'b [u8]) -> Option<&'b str> {
        if input.len() < 4 {
            return None;
        }
        let len = u32::from_be_bytes([input[0], input[1], input[2], input[3]]) as usize;
        *input = &input[4..];
        if input.len() < len {
            return None;
        }
        let str_bytes = &input[..len];
        *input = &input[len..];
        str::from_utf8(str_bytes).ok()
    }

    #[inline]
    pub fn put_hnsw_level_prefix(
        dst: &mut Vec<u8>,
        ns: &str,
        index: &str,
        field: &str,
        level: u16,
    ) {
        Self::put_namespace(dst, ns);
        Self::put_type(dst, SearchSubkeyType::Field);
        Self::put_sized_string(dst, index);
        Self::put_sized_string(dst, field);
        dst.extend_from_slice(&level.to_be_bytes());
    }

    #[inline]
    pub fn put_hnsw_level_node_prefix(
        dst: &mut Vec<u8>,
        ns: &str,
        index: &str,
        field: &str,
        level: u16,
    ) {
        Self::put_hnsw_level_prefix(dst, ns, index, field, level);
        dst.push(HnswLevelType::Node as u8);
    }

    #[inline]
    pub fn put_hnsw_level_edge_prefix(
        dst: &mut Vec<u8>,
        ns: &str,
        index: &str,
        field: &str,
        level: u16,
    ) {
        Self::put_hnsw_level_prefix(dst, ns, index, field, level);
        dst.push(HnswLevelType::Edge as u8);
    }

    #[inline]
    pub fn construct_index_meta(&self) -> Vec<u8> {
        let mut dst = Vec::with_capacity(1 + self.ns.len() + 1 + 4 + self.index.len());
        Self::put_namespace(&mut dst, self.ns);
        Self::put_type(&mut dst, SearchSubkeyType::IndexMeta);
        Self::put_sized_string(&mut dst, self.index);
        dst
    }

    #[inline]
    pub fn construct_index_prefixes(&self) -> Vec<u8> {
        let mut dst = Vec::with_capacity(1 + self.ns.len() + 1 + 4 + self.index.len());
        Self::put_namespace(&mut dst, self.ns);
        Self::put_type(&mut dst, SearchSubkeyType::Prefixes);
        Self::put_sized_string(&mut dst, self.index);
        dst
    }

    #[inline]
    pub fn construct_field_meta(&self) -> Vec<u8> {
        let field_name = self.field.unwrap_or("");
        let mut dst =
            Vec::with_capacity(1 + self.ns.len() + 1 + 4 + self.index.len() + 4 + field_name.len());
        Self::put_namespace(&mut dst, self.ns);
        Self::put_type(&mut dst, SearchSubkeyType::FieldMeta);
        Self::put_sized_string(&mut dst, self.index);
        Self::put_sized_string(&mut dst, field_name);
        dst
    }

    #[inline]
    pub fn construct_all_field_meta_begin(&self) -> Vec<u8> {
        let mut dst = Vec::with_capacity(1 + self.ns.len() + 1 + 4 + self.index.len() + 4);
        Self::put_namespace(&mut dst, self.ns);
        Self::put_type(&mut dst, SearchSubkeyType::FieldMeta);
        Self::put_sized_string(&mut dst, self.index);
        dst.extend_from_slice(&0u32.to_be_bytes());
        dst
    }

    #[inline]
    pub fn construct_all_field_meta_end(&self) -> Vec<u8> {
        let mut dst = Vec::with_capacity(1 + self.ns.len() + 1 + 4 + self.index.len() + 4);
        Self::put_namespace(&mut dst, self.ns);
        Self::put_type(&mut dst, SearchSubkeyType::FieldMeta);
        Self::put_sized_string(&mut dst, self.index);
        dst.extend_from_slice(&u32::MAX.to_be_bytes());
        dst
    }

    #[inline]
    pub fn construct_all_field_data_begin(&self) -> Vec<u8> {
        let mut dst = Vec::with_capacity(1 + self.ns.len() + 1 + 4 + self.index.len() + 4);
        Self::put_namespace(&mut dst, self.ns);
        Self::put_type(&mut dst, SearchSubkeyType::Field);
        Self::put_sized_string(&mut dst, self.index);
        dst.extend_from_slice(&0u32.to_be_bytes());
        dst
    }

    #[inline]
    pub fn construct_all_field_data_end(&self) -> Vec<u8> {
        let mut dst = Vec::with_capacity(1 + self.ns.len() + 1 + 4 + self.index.len() + 4);
        Self::put_namespace(&mut dst, self.ns);
        Self::put_type(&mut dst, SearchSubkeyType::Field);
        Self::put_sized_string(&mut dst, self.index);
        dst.extend_from_slice(&u32::MAX.to_be_bytes());
        dst
    }

    #[inline]
    pub fn construct_tag_field_data(&self, tag: &str, key: &str) -> Vec<u8> {
        let field_name = self.field.unwrap_or("");
        let mut dst = Vec::with_capacity(
            1 + self.ns.len()
                + 1
                + 4
                + self.index.len()
                + 4
                + field_name.len()
                + 4
                + tag.len()
                + 4
                + key.len(),
        );
        Self::put_namespace(&mut dst, self.ns);
        Self::put_type(&mut dst, SearchSubkeyType::Field);
        Self::put_sized_string(&mut dst, self.index);
        Self::put_sized_string(&mut dst, field_name);
        Self::put_sized_string(&mut dst, tag);
        Self::put_sized_string(&mut dst, key);
        dst
    }

    #[inline]
    pub fn construct_numeric_field_data(&self, num: f64, key: &str) -> Vec<u8> {
        let field_name = self.field.unwrap_or("");
        let mut dst = Vec::with_capacity(
            1 + self.ns.len() + 1 + 4 + self.index.len() + 4 + field_name.len() + 8 + 4 + key.len(),
        );
        Self::put_namespace(&mut dst, self.ns);
        Self::put_type(&mut dst, SearchSubkeyType::Field);
        Self::put_sized_string(&mut dst, self.index);
        Self::put_sized_string(&mut dst, field_name);
        dst.extend_from_slice(&encode_sortable_f64_u64(num).to_be_bytes());
        Self::put_sized_string(&mut dst, key);
        dst
    }

    #[inline]
    pub fn construct_hnsw_level_node_prefix(&self, level: u16) -> Vec<u8> {
        let field_name = self.field.unwrap_or("");
        let mut dst = Vec::new();
        Self::put_hnsw_level_node_prefix(&mut dst, self.ns, self.index, field_name, level);
        dst
    }

    #[inline]
    pub fn construct_hnsw_node(&self, level: u16, key: &str) -> Vec<u8> {
        let mut dst = self.construct_hnsw_level_node_prefix(level);
        Self::put_sized_string(&mut dst, key);
        dst
    }

    #[inline]
    pub fn construct_hnsw_edge_with_single_end(&self, level: u16, key: &str) -> Vec<u8> {
        let field_name = self.field.unwrap_or("");
        let mut dst = Vec::new();
        Self::put_hnsw_level_edge_prefix(&mut dst, self.ns, self.index, field_name, level);
        Self::put_sized_string(&mut dst, key);
        dst
    }

    #[inline]
    pub fn construct_hnsw_edge(&self, level: u16, key1: &str, key2: &str) -> Vec<u8> {
        let mut dst = self.construct_hnsw_edge_with_single_end(level, key1);
        Self::put_sized_string(&mut dst, key2);
        dst
    }
}

/// 索引元数据二进制编码（对标 Apache Kvrocks IndexMetadata::Encode）
#[inline]
pub fn encode_index_meta(on_data_type: IndexOnDataType) -> Vec<u8> {
    vec![0u8, on_data_type as u8]
}

/// 索引元数据二进制解码（对标 Apache Kvrocks IndexMetadata::Decode）
#[inline]
pub fn decode_index_meta(slice: &[u8]) -> Result<IndexOnDataType> {
    if slice.len() < 2 {
        return Err(Error::invalid_data(
            "insufficient length while decoding metadata",
        ));
    }
    match slice[1] {
        2 => Ok(IndexOnDataType::Hash),
        10 => Ok(IndexOnDataType::Json),
        other => Err(Error::invalid_data(format!(
            "unknown on_data_type: {other}"
        ))),
    }
}

/// 索引前缀列表二进制编码（对标 Apache Kvrocks IndexPrefixes::Encode）
pub fn encode_index_prefixes(prefixes: &[&str]) -> Vec<u8> {
    let mut dst = Vec::new();
    for prefix in prefixes {
        dst.extend_from_slice(&(prefix.len() as u32).to_be_bytes());
        dst.extend_from_slice(prefix.as_bytes());
    }
    dst
}

/// 索引前缀列表二进制解码（对标 Apache Kvrocks IndexPrefixes::Decode）
pub fn decode_index_prefixes(mut slice: &[u8]) -> Result<Vec<String>> {
    let mut prefixes = Vec::new();
    while slice.len() >= 4 {
        let len = u32::from_be_bytes([slice[0], slice[1], slice[2], slice[3]]) as usize;
        slice = &slice[4..];
        if slice.len() < len {
            return Err(Error::invalid_data(
                "insufficient length while decoding index prefixes",
            ));
        }
        let prefix_str = str::from_utf8(&slice[..len])
            .map_err(|_| Error::invalid_data("invalid utf-8 string in index prefixes"))?;
        prefixes.push(prefix_str.to_string());
        slice = &slice[len..];
    }
    Ok(prefixes)
}

/// 标签字段元数据二进制编码（对标 Apache Kvrocks TagFieldMetadata::Encode）
#[inline]
pub fn encode_tag_field_meta(separator: char, case_sensitive: bool, noindex: bool) -> Vec<u8> {
    let flag = (noindex as u8) | ((IndexFieldType::Tag as u8) << 1);
    vec![flag, separator as u8, case_sensitive as u8]
}

/// 标签字段元数据二进制解码（对标 Apache Kvrocks TagFieldMetadata::Decode）
#[inline]
pub fn decode_tag_field_meta(slice: &[u8]) -> Result<(char, bool, bool)> {
    if slice.len() < 3 {
        return Err(Error::invalid_data(
            "insufficient length while decoding tag field metadata",
        ));
    }
    let flag = slice[0];
    let noindex = (flag & 1) != 0;
    let separator = slice[1] as char;
    let case_sensitive = slice[2] != 0;
    Ok((separator, case_sensitive, noindex))
}

/// 数值字段元数据二进制编码（对标 Apache Kvrocks NumericFieldMetadata::Encode）
#[inline]
pub fn encode_numeric_field_meta(noindex: bool) -> Vec<u8> {
    let flag = (noindex as u8) | ((IndexFieldType::Numeric as u8) << 1);
    vec![flag]
}

/// 数值字段元数据二进制解码（对标 Apache Kvrocks NumericFieldMetadata::Decode）
#[inline]
pub fn decode_numeric_field_meta(slice: &[u8]) -> Result<bool> {
    if slice.is_empty() {
        return Err(Error::invalid_data(
            "insufficient length while decoding numeric field metadata",
        ));
    }
    let flag = slice[0];
    let noindex = (flag & 1) != 0;
    Ok(noindex)
}

/// HNSW 向量字段元数据编码字节长度 (1 flag + 1 type + 2 dim + 1 metric + 4 cap + 2 m + 4 ef_c + 4 ef_r + 8 epsilon + 2 num_levels)
pub const HNSW_VECTOR_FIELD_META_LEN: usize = 29;

/// HNSW 向量字段元数据二进制编码（对标 Apache Kvrocks HnswVectorFieldMetadata::Encode）
pub fn encode_hnsw_vector_field_meta(meta: &VectorFieldMetadata, noindex: bool) -> Vec<u8> {
    let flag = (noindex as u8) | ((IndexFieldType::Vector as u8) << 1);
    let mut dst = Vec::with_capacity(HNSW_VECTOR_FIELD_META_LEN);
    dst.push(flag);
    dst.push(meta.vector_type as u8);
    dst.extend_from_slice(&(meta.dim as u16).to_be_bytes());
    dst.push(meta.distance_metric as u8);
    dst.extend_from_slice(&(meta.initial_cap as u32).to_be_bytes());
    dst.extend_from_slice(&(meta.m as u16).to_be_bytes());
    dst.extend_from_slice(&(meta.ef_construction as u32).to_be_bytes());
    dst.extend_from_slice(&(meta.ef_runtime as u32).to_be_bytes());
    dst.extend_from_slice(&encode_sortable_f64_u64(meta.epsilon).to_be_bytes());
    dst.extend_from_slice(&meta.num_levels.to_be_bytes());
    dst
}

/// HNSW 向量字段元数据二进制解码（对标 Apache Kvrocks HnswVectorFieldMetadata::Decode）
pub fn decode_hnsw_vector_field_meta(slice: &[u8]) -> Result<(VectorFieldMetadata, bool)> {
    if slice.len() < HNSW_VECTOR_FIELD_META_LEN {
        return Err(Error::invalid_data(
            "insufficient length while decoding hnsw vector field metadata",
        ));
    }
    let flag = slice[0];
    let noindex = (flag & 1) != 0;

    let vector_type = match slice[1] {
        1 => VectorType::Float64,
        2 => VectorType::Float32,
        _ => VectorType::Float64,
    };
    let dim = u16::from_be_bytes([slice[2], slice[3]]) as usize;
    let distance_metric = match slice[4] {
        0 => DistanceMetric::L2,
        1 => DistanceMetric::IP,
        2 => DistanceMetric::Cosine,
        _ => DistanceMetric::Cosine,
    };
    let initial_cap = u32::from_be_bytes([slice[5], slice[6], slice[7], slice[8]]) as usize;
    let m = u16::from_be_bytes([slice[9], slice[10]]) as usize;
    let ef_construction = u32::from_be_bytes([slice[11], slice[12], slice[13], slice[14]]) as usize;
    let ef_runtime = u32::from_be_bytes([slice[15], slice[16], slice[17], slice[18]]) as usize;
    let epsilon_u64 = u64::from_be_bytes([
        slice[19], slice[20], slice[21], slice[22], slice[23], slice[24], slice[25], slice[26],
    ]);
    let epsilon = decode_sortable_f64_u64(epsilon_u64);
    let num_levels = u16::from_be_bytes([slice[27], slice[28]]);

    let meta = VectorFieldMetadata {
        vector_type,
        dim,
        distance_metric,
        algorithm: VectorAlgorithm::Hnsw,
        initial_cap,
        m,
        ef_construction,
        ef_runtime,
        epsilon,
        num_levels,
    };
    Ok((meta, noindex))
}

/// HNSW 节点元数据二进制编码（对标 Apache Kvrocks HnswNodeFieldMetadata::Encode）
pub fn encode_hnsw_node_meta(num_neighbours: u16, vector: &[f64]) -> Vec<u8> {
    let mut dst = Vec::with_capacity(4 + vector.len() * 8);
    dst.extend_from_slice(&num_neighbours.to_be_bytes());
    dst.extend_from_slice(&(vector.len() as u16).to_be_bytes());
    for &element in vector {
        dst.extend_from_slice(&encode_sortable_f64_u64(element).to_be_bytes());
    }
    dst
}

/// HNSW 节点元数据二进制解码（对标 Apache Kvrocks HnswNodeFieldMetadata::Decode）
pub fn decode_hnsw_node_meta(slice: &[u8]) -> Result<(u16, Vec<f64>)> {
    if slice.len() < 4 {
        return Err(Error::invalid_data(
            "insufficient length while decoding hnsw node metadata",
        ));
    }
    let num_neighbours = u16::from_be_bytes([slice[0], slice[1]]);
    let dim = u16::from_be_bytes([slice[2], slice[3]]) as usize;
    if slice.len() != 4 + dim * 8 {
        return Err(Error::invalid_data(
            "length is too short or too long to be parsed as a vector",
        ));
    }
    let mut vec = Vec::with_capacity(dim);
    for chunk in slice[4..].as_chunks::<8>().0 {
        let u = u64::from_be_bytes([
            chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], chunk[6], chunk[7],
        ]);
        vec.push(decode_sortable_f64_u64(u));
    }
    Ok((num_neighbours, vec))
}

/// 浮点数可排序 64 位整数转换（按位变换保证前缀扫描有序，IEEE 754 标准映射）
pub use crate::meta::{decode_sortable_f64_u64, encode_sortable_f64_u64};

/// 浮点数可排序十六进制字符串编码
#[inline]
pub fn encode_sortable_f64(val: f64) -> String {
    let encoded = encode_sortable_f64_u64(val);
    let bytes = u64_to_hex_16(encoded);
    unsafe { String::from_utf8_unchecked(bytes.to_vec()) }
}

/// 浮点数可排序十六进制字符串解码
#[inline]
pub fn decode_sortable_f64(hex_str: &str) -> Option<f64> {
    let sortable = decode_hex_u64(hex_str.as_bytes())?;
    Some(decode_sortable_f64_u64(sortable))
}

/// 有符号 64 位整数可排序十六进制编码
#[inline]
pub fn encode_sortable_i64(val: i64) -> String {
    let unsigned = (val as u64) ^ (1 << 63);
    let bytes = u64_to_hex_16(unsigned);
    unsafe { String::from_utf8_unchecked(bytes.to_vec()) }
}

/// 有符号 64 位整数可排序十六进制解码
#[inline]
pub fn decode_sortable_i64(hex_str: &str) -> Option<i64> {
    let unsigned = decode_hex_u64(hex_str.as_bytes())?;
    Some((unsigned ^ (1 << 63)) as i64)
}

/// 向量距离与相似度度量计算（对标 Apache Kvrocks ComputeSimilarity，单次迭代聚合）
#[inline]
pub fn compute_vector_distance(v1: &[f64], v2: &[f64], metric: DistanceMetric) -> Result<f64> {
    if v1.len() != v2.len() {
        let len1 = v1.len();
        let len2 = v2.len();
        return Err(Error::invalid_data(format!(
            "vector dimension mismatch: {len1} vs {len2}"
        )));
    }
    if v1.is_empty() {
        return Err(Error::invalid_data("empty vector is invalid"));
    }

    match metric {
        DistanceMetric::L2 => {
            let sum: f64 = v1.iter().zip(v2.iter()).fold(0.0, |acc, (&a, &b)| {
                let diff = a - b;
                acc + diff * diff
            });
            Ok(sum.sqrt())
        }
        DistanceMetric::IP => {
            let dot: f64 = v1
                .iter()
                .zip(v2.iter())
                .fold(0.0, |acc, (&a, &b)| acc + a * b);
            // 内积度量下，数值越大距离越小，故取负值使得升序排序与余弦/L2一致
            Ok(-dot)
        }
        DistanceMetric::Cosine => {
            let (dot, norm1, norm2) = v1.iter().zip(v2.iter()).fold(
                (0.0f64, 0.0f64, 0.0f64),
                |(dot_acc, n1_acc, n2_acc), (&a, &b)| {
                    (dot_acc + a * b, n1_acc + a * a, n2_acc + b * b)
                },
            );
            if norm1 <= 0.0 || norm2 <= 0.0 {
                return Ok(1.0);
            }
            let sim = dot / (norm1.sqrt() * norm2.sqrt());
            let sim_clamped = sim.clamp(-1.0, 1.0);
            Ok(1.0 - sim_clamped)
        }
    }
}

/// 从二进制字节数组或 JSON 数组中解析浮点向量（零冗余单次解析）
pub fn parse_vector_from_slice(bytes: &[u8], vector_type: VectorType) -> Result<Vec<f64>> {
    if bytes.is_empty() {
        return Err(Error::invalid_data("empty vector byte format"));
    }

    // 优先检测 JSON 格式（以 '[' 开头）
    if bytes.starts_with(b"[")
        && let Ok(json_v) = sonic_rs::from_slice::<sonic_rs::Value>(bytes)
        && let Some(arr) = json_v.as_array()
    {
        let mut vec = Vec::with_capacity(arr.len());
        for item in arr {
            if let Some(n) = item.as_f64() {
                vec.push(n);
            }
        }
        if !vec.is_empty() {
            return Ok(vec);
        }
    }

    // 二进制字节数组解析 (Float64: 8 字节/元素, Float32: 4 字节/元素)
    let elem_size = vector_type.byte_size();
    if bytes.len().is_multiple_of(elem_size) {
        match vector_type {
            VectorType::Float64 => {
                let count = bytes.len() / 8;
                let mut vec = Vec::with_capacity(count);
                for chunk in bytes.as_chunks::<8>().0 {
                    vec.push(f64::from_le_bytes(*chunk));
                }
                return Ok(vec);
            }
            VectorType::Float32 => {
                let count = bytes.len() / 4;
                let mut vec = Vec::with_capacity(count);
                for chunk in bytes.as_chunks::<4>().0 {
                    vec.push(f32::from_le_bytes(*chunk) as f64);
                }
                return Ok(vec);
            }
        }
    }

    // 纯文本逗号分隔（含 ',' 且为有效 UTF-8 文本）
    if let Ok(s) = str::from_utf8(bytes)
        && s.contains(',')
    {
        let clean = s.trim().trim_start_matches('[').trim_end_matches(']');
        let mut vec = Vec::new();
        let mut valid = true;
        for part in clean.split(',') {
            let p = part.trim();
            if p.is_empty() {
                continue;
            }
            if let Ok(num) = p.parse::<f64>() {
                vec.push(num);
            } else {
                valid = false;
                break;
            }
        }
        if valid && !vec.is_empty() {
            return Ok(vec);
        }
    }

    Err(Error::invalid_data("invalid vector byte format or length"))
}