wedb_embed 0.1.1

use crate::search::conf::SuggestionItem;
use crate::search::tokenizer::levenshtein_distance;
use rapidhash::RapidHashMap as HashMap;
use std::cmp::Ordering;

/// 默认建议项检索限制数量
pub const DEFAULT_SUG_LIMIT: usize = 10;

/// 自动补全建议字典（对标 RediSearch FT.SUGADD, FT.SUGGET, FT.SUGDEL, FT.SUGLEN）
#[derive(Debug, Clone, Default)]
pub struct SuggestionDict {
    pub entries: HashMap<String, (f64, Option<String>)>,
}

impl SuggestionDict {
    #[inline]
    pub fn new() -> Self {
        Self::default()
    }

    /// 添加建议项
    pub fn sug_add(
        &mut self,
        string: &str,
        score: f64,
        incr: bool,
        payload: Option<String>,
    ) -> usize {
        let entry = self
            .entries
            .entry(string.to_string())
            .or_insert((0.0, None));
        if incr {
            entry.0 += score;
        } else {
            entry.0 = score;
        }
        if payload.is_some() {
            entry.1 = payload;
        }
        self.entries.len()
    }

    /// 检索建议项（快速前缀与编辑距离过滤 + 依据真实权重降序截断排序）
    pub fn sug_get(
        &self,
        prefix: &str,
        fuzzy: bool,
        withscores: bool,
        withpayloads: bool,
        max: Option<usize>,
    ) -> Vec<SuggestionItem> {
        let limit = max.unwrap_or(DEFAULT_SUG_LIMIT);
        let prefix_lower = prefix.to_lowercase();

        // 收集符合前缀或编辑距离的条目：(string, real_score, payload)
        let mut matched: Vec<(&str, f64, Option<&str>)> = self
            .entries
            .iter()
            .filter_map(|(s, (score, payload))| {
                let s_str = s.as_str();
                // 快速 ASCII 前缀检查
                if s_str.len() >= prefix.len() && s_str[..prefix.len()].eq_ignore_ascii_case(prefix)
                {
                    Some((s_str, *score, payload.as_deref()))
                } else {
                    let s_lower = s.to_lowercase();
                    if s_lower.starts_with(&prefix_lower)
                        || (fuzzy && levenshtein_distance(&s_lower, &prefix_lower) <= 1)
                    {
                        Some((s_str, *score, payload.as_deref()))
                    } else {
                        None
                    }
                }
            })
            .collect();

        // 依据真实分数降序排序，分数相同按字母序升序
        let cmp_fn = |a: &(&str, f64, Option<&str>), b: &(&str, f64, Option<&str>)| {
            b.1.partial_cmp(&a.1)
                .unwrap_or(Ordering::Equal)
                .then_with(|| a.0.cmp(b.0))
        };

        if matched.len() > limit {
            matched.select_nth_unstable_by(limit, cmp_fn);
            matched.truncate(limit);
        }

        matched.sort_by(cmp_fn);

        // 格式化输出 SuggestionItem
        matched
            .into_iter()
            .map(|(s, score, payload)| SuggestionItem {
                string: s.to_string(),
                score: if withscores { score } else { 0.0 },
                payload: if withpayloads {
                    payload.map(str::to_string)
                } else {
                    None
                },
            })
            .collect()
    }

    /// 删除建议项
    #[inline]
    pub fn sug_del(&mut self, string: &str) -> bool {
        self.entries.remove(string).is_some()
    }

    /// 返回建议字典长度
    #[inline]
    pub fn sug_len(&self) -> usize {
        self.entries.len()
    }
}