wedb_embed 0.1.0

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

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

impl SuggestionDict {
    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(10);
        let prefix_lower = prefix.to_lowercase();

        let mut matched: Vec<SuggestionItem> = self
            .entries
            .iter()
            .filter(|(s, _)| {
                let s_lower = s.to_lowercase();
                if s_lower.starts_with(&prefix_lower) {
                    true
                } else if fuzzy {
                    levenshtein_distance(&s_lower, &prefix_lower) <= 1
                } else {
                    false
                }
            })
            .map(|(s, (score, payload))| SuggestionItem {
                string: s.clone(),
                score: if withscores { *score } else { 0.0 },
                payload: if withpayloads { payload.clone() } else { None },
            })
            .collect();

        if matched.len() > limit {
            matched.select_nth_unstable_by(limit, |a, b| {
                b.score
                    .partial_cmp(&a.score)
                    .unwrap_or(Ordering::Equal)
                    .then_with(|| a.string.cmp(&b.string))
            });
            matched.truncate(limit);
        }

        // 依据分数由大到小排序，分数相同按字母序排序
        matched.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(Ordering::Equal)
                .then_with(|| a.string.cmp(&b.string))
        });

        matched
    }

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

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