wedb_embed 0.1.1

use crate::error::{Error, Result};
use crate::search::encoding::compute_vector_distance;
use crate::search::meta::DistanceMetric;
use hipstr::HipStr;
use rapidhash::{RapidHashMap as HashMap, RapidHashSet as HashSet};
use std::cmp::Ordering;
use std::collections::BinaryHeap;

/// 单个候选向量距离对（用于最小堆/最大堆）
#[derive(Debug, Clone)]
pub struct Candidate {
    pub dist: f64,
    pub doc_id: HipStr<'static>,
}

impl PartialEq for Candidate {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        self.dist == other.dist && self.doc_id == other.doc_id
    }
}

impl Eq for Candidate {}

impl PartialOrd for Candidate {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Candidate {
    #[inline]
    fn cmp(&self, other: &Self) -> Ordering {
        // 用于默认 BinaryHeap（大顶堆：距离大的在堆顶）
        self.dist
            .partial_cmp(&other.dist)
            .unwrap_or(Ordering::Equal)
            .then_with(|| self.doc_id.cmp(&other.doc_id))
    }
}

/// 逆序候选者（用于小顶堆：距离小的在堆顶）
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MinCandidate(pub Candidate);

impl PartialOrd for MinCandidate {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for MinCandidate {
    #[inline]
    fn cmp(&self, other: &Self) -> Ordering {
        other.0.cmp(&self.0)
    }
}

/// HNSW 图节点（对标 Apache Kvrocks HnswNode）
#[derive(Debug, Clone)]
pub struct HnswNode {
    pub doc_id: HipStr<'static>,
    pub vector: Vec<f64>,
    pub level: usize,
    /// 每一层的邻居列表，索引为层号 0..=level
    pub neighbors: Vec<Vec<HipStr<'static>>>,
}

/// 内存 HNSW 向量索引图（对标 Apache Kvrocks HnswIndex 与 RediSearch HNSW 实现）
#[derive(Debug, Clone)]
pub struct HnswGraph {
    pub dim: usize,
    pub distance_metric: DistanceMetric,
    pub m: usize,
    pub ef_construction: usize,
    pub ef_runtime: usize,
    pub epsilon: f64,
    pub max_level: usize,
    pub entry_point: Option<HipStr<'static>>,
    pub nodes: HashMap<HipStr<'static>, HnswNode>,
    level_mult: f64,
}

impl Default for HnswGraph {
    fn default() -> Self {
        Self::new(0, DistanceMetric::Cosine, 16, 200, 10, 0.01)
    }
}

impl HnswGraph {
    pub fn new(
        dim: usize,
        distance_metric: DistanceMetric,
        m: usize,
        ef_construction: usize,
        ef_runtime: usize,
        epsilon: f64,
    ) -> Self {
        let m_val = m.max(2);
        let level_mult = 1.0 / (m_val as f64).ln();
        Self {
            dim,
            distance_metric,
            m: m_val,
            ef_construction: ef_construction.max(1),
            ef_runtime: ef_runtime.max(1),
            epsilon,
            max_level: 0,
            entry_point: None,
            nodes: HashMap::default(),
            level_mult,
        }
    }

    /// 返回当前图的层级数（对标 Apache Kvrocks metadata.num_levels）
    #[inline]
    pub fn num_levels(&self) -> u16 {
        if self.nodes.is_empty() {
            0
        } else {
            (self.max_level + 1) as u16
        }
    }

    /// 随机生成新插入节点的层数（对标 Apache Kvrocks HnswIndex::RandomizeLayer）
    #[inline]
    pub fn random_level(&self) -> usize {
        let r: f64 = fastrand::f64();
        let r = r.max(f64::MIN_POSITIVE);
        ((-r.ln()) * self.level_mult).floor() as usize
    }

    /// 计算两向量间距离
    #[inline]
    pub fn dist(&self, v1: &[f64], v2: &[f64]) -> Result<f64> {
        compute_vector_distance(v1, v2, self.distance_metric)
    }

    /// 插入新向量（对标 Apache Kvrocks HnswIndex::InsertVectorEntry）
    pub fn insert(&mut self, doc_id: HipStr<'static>, vector: Vec<f64>) -> Result<()> {
        if self.dim != 0 && vector.len() != self.dim {
            let dim = self.dim;
            let len = vector.len();
            return Err(Error::invalid_data(format!(
                "vector dimension mismatch: expected {dim}, got {len}"
            )));
        }
        if self.dim == 0 {
            self.dim = vector.len();
        }

        // 如果节点已存在，先删除旧节点
        if self.nodes.contains_key(&doc_id) {
            self.delete(doc_id.as_str());
        }

        let node_level = self.random_level();
        let neighbors = vec![Vec::new(); node_level + 1];

        let new_node = HnswNode {
            doc_id: doc_id.clone(),
            vector: vector.clone(),
            level: node_level,
            neighbors,
        };
        self.nodes.insert(doc_id.clone(), new_node);

        let Some(mut curr_ep) = self.entry_point.clone() else {
            self.entry_point = Some(doc_id);
            self.max_level = node_level;
            return Ok(());
        };
        let max_lvl = self.max_level;

        // 1. 从最高层下行到 node_level + 1 贪心搜索最近入口点
        if max_lvl > node_level {
            for lvl in (node_level + 1..=max_lvl).rev() {
                let mut changed = true;
                while changed {
                    changed = false;
                    let ep_node = match self.nodes.get(&curr_ep) {
                        Some(n) => n,
                        None => break,
                    };
                    let curr_dist = self.dist(&vector, &ep_node.vector)?;
                    let mut closest_ep = curr_ep.clone();
                    let mut min_d = curr_dist;

                    if lvl < ep_node.neighbors.len() {
                        for neighbor_id in &ep_node.neighbors[lvl] {
                            if let Some(neighbor_node) = self.nodes.get(neighbor_id) {
                                let d = self.dist(&vector, &neighbor_node.vector)?;
                                if d < min_d {
                                    min_d = d;
                                    closest_ep = neighbor_id.clone();
                                    changed = true;
                                }
                            }
                        }
                    }
                    if changed {
                        curr_ep = closest_ep;
                    }
                }
            }
        }

        // 2. 从 min(node_level, max_lvl) 到 0 层逐层连边
        let mut curr_eps = vec![curr_ep];
        let insert_top_level = node_level.min(max_lvl);

        for lvl in (0..=insert_top_level).rev() {
            let candidates =
                self.search_layer_internal(&vector, &curr_eps, self.ef_construction, lvl)?;
            let m_max = if lvl == 0 { self.m * 2 } else { self.m };
            let selected = self.select_neighbors(&candidates, m_max);

            // 连接新节点与选取的邻居
            if let Some(node) = self.nodes.get_mut(&doc_id) {
                node.neighbors[lvl] = selected.clone();
            }

            // 双向连边与邻居截断
            for neighbor_id in &selected {
                if let Some(neighbor_node) = self.nodes.get_mut(neighbor_id)
                    && lvl < neighbor_node.neighbors.len()
                {
                    if !neighbor_node.neighbors[lvl].contains(&doc_id) {
                        neighbor_node.neighbors[lvl].push(doc_id.clone());
                    }
                    if neighbor_node.neighbors[lvl].len() > m_max {
                        let n_vec = neighbor_node.vector.clone();
                        let n_candidates: Vec<HipStr<'static>> =
                            neighbor_node.neighbors[lvl].clone();
                        let pruned = self.select_neighbors_from_ids(&n_vec, &n_candidates, m_max);
                        if let Some(n_node_re) = self.nodes.get_mut(neighbor_id) {
                            n_node_re.neighbors[lvl] = pruned;
                        }
                    }
                }
            }

            curr_eps = candidates.into_iter().map(|c| c.doc_id).collect();
        }

        if node_level > self.max_level {
            self.max_level = node_level;
            self.entry_point = Some(doc_id);
        }

        Ok(())
    }

    /// 删除节点（对标 Apache Kvrocks HnswIndex::DeleteVectorEntry）
    pub fn delete(&mut self, doc_id: &str) -> bool {
        let doc_key = HipStr::from(doc_id);
        let removed = match self.nodes.remove(&doc_key) {
            Some(n) => n,
            None => return false,
        };

        // 清理所有邻居的反向连边
        for (lvl, n_list) in removed.neighbors.iter().enumerate() {
            for neighbor_id in n_list {
                if let Some(neighbor_node) = self.nodes.get_mut(neighbor_id)
                    && lvl < neighbor_node.neighbors.len()
                {
                    neighbor_node.neighbors[lvl].retain(|id| id.as_str() != doc_id);
                }
            }
        }

        // 若删除的是入口点，更新入口点为剩余节点中层数最高者
        if self.entry_point.as_deref() == Some(doc_id) {
            if self.nodes.is_empty() {
                self.entry_point = None;
                self.max_level = 0;
            } else {
                let mut best_ep = None;
                let mut best_lvl = 0;
                for (id, node) in &self.nodes {
                    if best_ep.is_none() || node.level >= best_lvl {
                        best_ep = Some(id.clone());
                        best_lvl = node.level;
                    }
                }
                self.entry_point = best_ep;
                self.max_level = best_lvl;
            }
        }

        true
    }

    /// 单层 Beam Search 搜索邻近候选集（对标 Apache Kvrocks HnswIndex::SearchLayerInternal）
    pub fn search_layer_internal(
        &self,
        query: &[f64],
        entry_points: &[HipStr<'static>],
        ef: usize,
        level: usize,
    ) -> Result<Vec<Candidate>> {
        let mut visited = HashSet::default();
        let mut explore_heap = BinaryHeap::new(); // 小顶堆 MinCandidate
        let mut result_heap = BinaryHeap::new(); // 大顶堆 Candidate (保持最近的 ef 个)

        for ep in entry_points {
            if let Some(ep_node) = self.nodes.get(ep) {
                let dist = self.dist(query, &ep_node.vector)?;
                let cand = Candidate {
                    dist,
                    doc_id: ep.clone(),
                };
                explore_heap.push(MinCandidate(cand.clone()));
                result_heap.push(cand);
                visited.insert(ep.clone());
            }
        }

        while let Some(MinCandidate(curr)) = explore_heap.pop() {
            if let Some(furthest) = result_heap.peek()
                && curr.dist > furthest.dist
            {
                break;
            }

            if let Some(curr_node) = self.nodes.get(&curr.doc_id)
                && level < curr_node.neighbors.len()
            {
                for neighbor_id in &curr_node.neighbors[level] {
                    if !visited.insert(neighbor_id.clone()) {
                        continue;
                    }
                    if let Some(neighbor_node) = self.nodes.get(neighbor_id) {
                        let dist = self.dist(query, &neighbor_node.vector)?;
                        let cand = Candidate {
                            dist,
                            doc_id: neighbor_id.clone(),
                        };

                        if result_heap.len() < ef
                            || dist < result_heap.peek().map(|f| f.dist).unwrap_or(f64::INFINITY)
                        {
                            explore_heap.push(MinCandidate(cand.clone()));
                            result_heap.push(cand);
                            if result_heap.len() > ef {
                                result_heap.pop();
                            }
                        }
                    }
                }
            }
        }

        let mut res: Vec<Candidate> = result_heap.into_vec();
        res.sort();
        Ok(res)
    }

    /// 邻居启发式筛选（对标 Apache Kvrocks HnswIndex::SelectNeighbors）
    #[inline]
    pub fn select_neighbors(&self, candidates: &[Candidate], m_max: usize) -> Vec<HipStr<'static>> {
        candidates
            .iter()
            .take(m_max)
            .map(|c| c.doc_id.clone())
            .collect()
    }

    pub fn select_neighbors_from_ids(
        &self,
        base_vec: &[f64],
        candidates: &[HipStr<'static>],
        m_max: usize,
    ) -> Vec<HipStr<'static>> {
        let mut scored: Vec<(f64, HipStr<'static>)> = Vec::with_capacity(candidates.len());
        for id in candidates {
            if let Some(node) = self.nodes.get(id)
                && let Ok(d) = self.dist(base_vec, &node.vector)
            {
                scored.push((d, id.clone()));
            }
        }
        scored.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal));
        scored.into_iter().take(m_max).map(|(_, id)| id).collect()
    }

    /// 执行 KNN 近邻检索（对标 Apache Kvrocks HnswIndex::KnnSearch）
    pub fn search_knn(
        &self,
        query: &[f64],
        k: usize,
        ef_runtime: Option<usize>,
    ) -> Result<Vec<(f64, HipStr<'static>)>> {
        let Some(mut curr_ep) = self.entry_point.clone() else {
            return Ok(Vec::new());
        };
        if self.nodes.is_empty() {
            return Ok(Vec::new());
        }

        let max_lvl = self.max_level;

        // 1. 上层贪心搜索入口点
        for lvl in (1..=max_lvl).rev() {
            let mut changed = true;
            while changed {
                changed = false;
                let ep_node = match self.nodes.get(&curr_ep) {
                    Some(n) => n,
                    None => break,
                };
                let curr_dist = self.dist(query, &ep_node.vector)?;
                let mut closest_ep = curr_ep.clone();
                let mut min_d = curr_dist;

                if lvl < ep_node.neighbors.len() {
                    for neighbor_id in &ep_node.neighbors[lvl] {
                        if let Some(neighbor_node) = self.nodes.get(neighbor_id) {
                            let d = self.dist(query, &neighbor_node.vector)?;
                            if d < min_d {
                                min_d = d;
                                closest_ep = neighbor_id.clone();
                                changed = true;
                            }
                        }
                    }
                }
                if changed {
                    curr_ep = closest_ep;
                }
            }
        }

        // 2. 底层 0 层 Beam Search
        let ef = ef_runtime.unwrap_or(self.ef_runtime).max(k);
        let candidates = self.search_layer_internal(query, &[curr_ep], ef, 0)?;

        let results = candidates
            .into_iter()
            .take(k)
            .map(|c| (c.dist, c.doc_id))
            .collect();
        Ok(results)
    }

    /// 执行范围检索（对标 Apache Kvrocks VECTOR_RANGE 检索）
    pub fn search_range(
        &self,
        query: &[f64],
        radius: f64,
        epsilon: Option<f64>,
    ) -> Result<Vec<(f64, HipStr<'static>)>> {
        let eps = epsilon.unwrap_or(self.epsilon);
        let effective_radius = radius * (1.0 + eps);
        let knn_candidates = self.search_knn(query, self.nodes.len(), Some(self.ef_runtime * 2))?;
        let filtered: Vec<(f64, HipStr<'static>)> = knn_candidates
            .into_iter()
            .filter(|(d, _)| *d <= effective_radius)
            .collect();
        Ok(filtered)
    }

    /// 扩展搜索范围（对标 Apache Kvrocks HnswIndex::ExpandSearchScope）
    pub fn expand_search_scope(
        &self,
        query: &[f64],
        initial_keys: &[(f64, HipStr<'static>)],
        visited: &mut HashSet<HipStr<'static>>,
    ) -> Result<Vec<(f64, HipStr<'static>)>> {
        let mut result = Vec::new();
        for (_, key) in initial_keys {
            if let Some(node) = self.nodes.get(key)
                && !node.neighbors.is_empty()
            {
                for neighbor_id in &node.neighbors[0] {
                    if !visited.insert(neighbor_id.clone()) {
                        continue;
                    }
                    if let Some(neighbor_node) = self.nodes.get(neighbor_id) {
                        let dist = self.dist(query, &neighbor_node.vector)?;
                        result.push((dist, neighbor_id.clone()));
                    }
                }
            }
        }
        result.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal));
        Ok(result)
    }

    /// 清空索引
    #[inline]
    pub fn clear(&mut self) {
        self.nodes.clear();
        self.entry_point = None;
        self.max_level = 0;
    }
}