//! HNSW (Hierarchical Navigable Small World) graph for approximate nearest
//! neighbor search.
//!
//! See [[hierarchical-navigable-small-world]] and [[architecture-overview#Step 2]].

use std::collections::{BinaryHeap, HashSet};
use std::cmp::Ordering;

use super::DistanceMetric;

/// HNSW index parameters.
#[derive(Clone, Debug)]
pub struct HnswParams {
    pub dims: usize,
    pub m: usize,              // max connections per node per layer
    pub ef_construction: usize, // beam width during construction
    pub metric: DistanceMetric,
}

impl Default for HnswParams {
    fn default() -> Self {
        Self {
            dims: 128,
            m: 16,
            ef_construction: 100,
            metric: DistanceMetric::Cosine,
        }
    }
}

/// A scored neighbor candidate (lower distance = better).
#[derive(Clone, Copy)]
struct Candidate {
    id: u32,
    dist: f32,
}

impl PartialEq for Candidate { fn eq(&self, o: &Self) -> bool { self.dist == o.dist } }
impl Eq for Candidate {}
impl PartialOrd for Candidate { fn partial_cmp(&self, o: &Self) -> Option<Ordering> { Some(self.cmp(o)) } }
// Min-heap: lower distance first
impl Ord for Candidate {
    fn cmp(&self, o: &Self) -> Ordering {
        o.dist.partial_cmp(&self.dist).unwrap_or(Ordering::Equal)
    }
}

/// Max-heap version for finding furthest candidates.
struct FurthestCandidate(Candidate);
impl PartialEq for FurthestCandidate { fn eq(&self, o: &Self) -> bool { self.0.dist == o.0.dist } }
impl Eq for FurthestCandidate {}
impl PartialOrd for FurthestCandidate { fn partial_cmp(&self, o: &Self) -> Option<Ordering> { Some(self.cmp(o)) } }
impl Ord for FurthestCandidate {
    fn cmp(&self, o: &Self) -> Ordering {
        self.0.dist.partial_cmp(&o.0.dist).unwrap_or(Ordering::Equal)
    }
}

/// Node in the HNSW graph.
struct Node {
    /// Connections per layer: layer_index → vec of neighbor ids
    neighbors: Vec<Vec<u32>>,
    level: usize,
}

/// HNSW graph builder.
pub struct HnswBuilder {
    params: HnswParams,
    vectors: Vec<Vec<f32>>,
    nodes: Vec<Node>,
    entry_point: Option<u32>,
    max_level: usize,
    level_mult: f64,
    rng_state: u64,
}

impl HnswBuilder {
    pub fn new(params: HnswParams) -> Self {
        let level_mult = 1.0 / (params.m as f64).ln();
        Self {
            params,
            vectors: Vec::new(),
            nodes: Vec::new(),
            entry_point: None,
            max_level: 0,
            level_mult,
            rng_state: 42,
        }
    }

    fn next_rand(&mut self) -> f64 {
        // Simple xorshift64 for deterministic tests
        self.rng_state ^= self.rng_state << 13;
        self.rng_state ^= self.rng_state >> 7;
        self.rng_state ^= self.rng_state << 17;
        (self.rng_state as f64) / (u64::MAX as f64)
    }

    fn random_level(&mut self) -> usize {
        let r = self.next_rand().max(1e-10);
        (-r.ln() * self.level_mult).floor() as usize
    }

    /// Add a vector to the index.
    pub fn add(&mut self, vector: Vec<f32>) {
        debug_assert_eq!(vector.len(), self.params.dims);
        let id = self.vectors.len() as u32;
        let level = self.random_level();

        self.vectors.push(vector);
        let mut neighbors = Vec::with_capacity(level + 1);
        for _ in 0..=level {
            neighbors.push(Vec::new());
        }
        self.nodes.push(Node { neighbors, level });

        if self.entry_point.is_none() {
            self.entry_point = Some(id);
            self.max_level = level;
            return;
        }

        let ep = self.entry_point.unwrap();

        // Greedy search from top to the insertion level + 1
        let mut current = ep;
        for lev in (level + 1..=self.max_level).rev() {
            current = self.greedy_closest(current, id, lev);
        }

        // Insert at each level from min(level, max_level) down to 0
        let insert_from = level.min(self.max_level);
        let mut ep_for_level = current;

        for lev in (0..=insert_from).rev() {
            let candidates = self.search_layer(id, ep_for_level, self.params.ef_construction, lev);
            let neighbors_to_connect = self.select_neighbors(&candidates, self.params.m);

            // Set this node's neighbors at this level
            if lev < self.nodes[id as usize].neighbors.len() {
                self.nodes[id as usize].neighbors[lev] = neighbors_to_connect.iter().map(|c| c.id).collect();
            }

            // Add back-links
            for &n in &neighbors_to_connect {
                if (lev) < self.nodes[n.id as usize].neighbors.len() {
                    self.nodes[n.id as usize].neighbors[lev].push(id);
                    // Prune if too many connections
                    if self.nodes[n.id as usize].neighbors[lev].len() > self.params.m * 2 {
                        self.prune_connections(n.id, lev);
                    }
                }
            }

            if !candidates.is_empty() {
                ep_for_level = candidates[0].id;
            }
        }

        if level > self.max_level {
            self.max_level = level;
            self.entry_point = Some(id);
        }
    }

    fn dist(&self, a: u32, b: u32) -> f32 {
        super::distance(&self.vectors[a as usize], &self.vectors[b as usize], self.params.metric)
    }

    fn dist_to_vec(&self, a: u32, query: &[f32]) -> f32 {
        super::distance(&self.vectors[a as usize], query, self.params.metric)
    }

    fn greedy_closest(&self, start: u32, target: u32, level: usize) -> u32 {
        let mut current = start;
        let mut current_dist = self.dist(current, target);
        loop {
            let mut changed = false;
            if level < self.nodes[current as usize].neighbors.len() {
                for &neighbor in &self.nodes[current as usize].neighbors[level] {
                    let d = self.dist(neighbor, target);
                    if d < current_dist {
                        current = neighbor;
                        current_dist = d;
                        changed = true;
                    }
                }
            }
            if !changed { break; }
        }
        current
    }

    fn search_layer(&self, query_id: u32, entry: u32, ef: usize, level: usize) -> Vec<Candidate> {
        self.search_layer_vec(&self.vectors[query_id as usize], entry, ef, level, None)
    }

    fn search_layer_vec(
        &self, query: &[f32], entry: u32, ef: usize, level: usize,
        filter: Option<&roaring::RoaringBitmap>,
    ) -> Vec<Candidate> {
        let mut visited = HashSet::new();
        let mut candidates = BinaryHeap::new(); // min-heap (closest first)
        let mut results = BinaryHeap::new(); // max-heap (furthest first)

        let d = self.dist_to_vec(entry, query);
        visited.insert(entry);
        candidates.push(Candidate { id: entry, dist: d });
        if filter.is_none() || filter.unwrap().contains(entry) {
            results.push(FurthestCandidate(Candidate { id: entry, dist: d }));
        }

        while let Some(c) = candidates.pop() {
            let furthest_dist = results.peek().map(|f| f.0.dist).unwrap_or(f32::MAX);
            if c.dist > furthest_dist && results.len() >= ef {
                break;
            }

            if level < self.nodes[c.id as usize].neighbors.len() {
                for &neighbor in &self.nodes[c.id as usize].neighbors[level] {
                    if visited.insert(neighbor) {
                        let d = self.dist_to_vec(neighbor, query);
                        let furthest_dist = results.peek().map(|f| f.0.dist).unwrap_or(f32::MAX);
                        if d < furthest_dist || results.len() < ef {
                            candidates.push(Candidate { id: neighbor, dist: d });
                            if filter.is_none() || filter.unwrap().contains(neighbor) {
                                results.push(FurthestCandidate(Candidate { id: neighbor, dist: d }));
                                if results.len() > ef {
                                    results.pop();
                                }
                            }
                        }
                    }
                }
            }
        }

        results.into_sorted_vec().into_iter().map(|f| f.0).collect()
    }

    fn select_neighbors(&self, candidates: &[Candidate], m: usize) -> Vec<Candidate> {
        let mut sorted = candidates.to_vec();
        sorted.sort_by(|a, b| a.dist.partial_cmp(&b.dist).unwrap_or(Ordering::Equal));
        sorted.truncate(m);
        sorted
    }

    fn prune_connections(&mut self, node_id: u32, level: usize) {
        let neighbors = &self.nodes[node_id as usize].neighbors[level];
        let mut candidates: Vec<Candidate> = neighbors.iter()
            .map(|&n| Candidate { id: n, dist: self.dist(node_id, n) })
            .collect();
        candidates.sort_by(|a, b| a.dist.partial_cmp(&b.dist).unwrap_or(Ordering::Equal));
        candidates.truncate(self.params.m);
        self.nodes[node_id as usize].neighbors[level] = candidates.iter().map(|c| c.id).collect();
    }

    /// Number of vectors in the index.
    pub fn len(&self) -> usize { self.vectors.len() }
    pub fn is_empty(&self) -> bool { self.vectors.is_empty() }

    /// Build the final index.
    pub fn build(self) -> HnswIndex {
        // Quantize vectors for fast approximate search
        let quantized = if !self.vectors.is_empty() {
            Some(super::quantize::QuantizedVectors::quantize(
                &self.vectors, self.params.metric,
            ))
        } else {
            None
        };

        HnswIndex {
            params: self.params,
            vectors: self.vectors,
            nodes: self.nodes,
            entry_point: self.entry_point,
            max_level: self.max_level,
            quantized,
            cached_query_norm: std::cell::Cell::new(0.0),
        }
    }
}

/// Immutable HNSW index for search.
pub struct HnswIndex {
    params: HnswParams,
    vectors: Vec<Vec<f32>>,
    nodes: Vec<Node>,
    entry_point: Option<u32>,
    max_level: usize,
    /// Int8 quantized vectors for fast approximate distance.
    quantized: Option<super::quantize::QuantizedVectors>,
    /// Precomputed query norm cache (for cosine distance).
    cached_query_norm: std::cell::Cell<f32>,
}

impl HnswIndex {
    /// Search for k nearest neighbors.
    pub fn search(&self, query: &[f32], k: usize, ef: usize) -> Vec<(u32, f32)> {
        self.search_filtered(query, k, ef, None)
    }

    /// Search with optional pre-filter bitmap.
    pub fn search_filtered(
        &self, query: &[f32], k: usize, ef: usize,
        filter: Option<&roaring::RoaringBitmap>,
    ) -> Vec<(u32, f32)> {
        if self.entry_point.is_none() {
            return Vec::new();
        }

        // Cache query norm for cosine distance
        if self.params.metric == super::DistanceMetric::Cosine {
            self.cached_query_norm.set(super::quantize::QuantizedVectors::query_norm(query));
        }

        // Brute-force fallback for very selective filters
        if let Some(bm) = filter {
            if (bm.len() as f64) < (self.vectors.len() as f64 * 0.01) {
                return self.brute_force_search(query, k, bm);
            }
        }

        let ep = self.entry_point.unwrap();
        let ef_actual = ef.max(k);

        // Greedy descend from top to layer 1
        let mut current = ep;
        for lev in (1..=self.max_level).rev() {
            current = self.greedy_closest_vec(current, query, lev);
        }

        // Search layer 0 with beam (uses quantized distance if available)
        let candidates = self.search_layer_0(query, current, ef_actual, filter);

        // Rerank with exact float32 distances for final ordering
        let mut results: Vec<(u32, f32)> = if self.quantized.is_some() {
            candidates.into_iter()
                .map(|c| (c.id, self.exact_dist(c.id, query)))
                .collect()
        } else {
            candidates.into_iter()
                .map(|c| (c.id, c.dist))
                .collect()
        };
        results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
        results.truncate(k);
        results
    }

    fn greedy_closest_vec(&self, start: u32, query: &[f32], level: usize) -> u32 {
        let mut current = start;
        let mut current_dist = self.quantized_dist(current, query);
        loop {
            let mut changed = false;
            if level < self.nodes[current as usize].neighbors.len() {
                for &neighbor in &self.nodes[current as usize].neighbors[level] {
                    let d = self.quantized_dist(neighbor, query);
                    if d < current_dist {
                        current = neighbor;
                        current_dist = d;
                        changed = true;
                    }
                }
            }
            if !changed { break; }
        }
        current
    }

    fn search_layer_0(
        &self, query: &[f32], entry: u32, ef: usize,
        filter: Option<&roaring::RoaringBitmap>,
    ) -> Vec<Candidate> {
        let mut visited = HashSet::new();
        let mut candidates = BinaryHeap::new();
        let mut results = BinaryHeap::new();

        let d = self.quantized_dist(entry, query);
        visited.insert(entry);
        candidates.push(Candidate { id: entry, dist: d });
        if filter.is_none() || filter.unwrap().contains(entry) {
            results.push(FurthestCandidate(Candidate { id: entry, dist: d }));
        }

        while let Some(c) = candidates.pop() {
            let furthest_dist = results.peek().map(|f| f.0.dist).unwrap_or(f32::MAX);
            if c.dist > furthest_dist && results.len() >= ef {
                break;
            }

            if !self.nodes[c.id as usize].neighbors.is_empty() {
                for &neighbor in &self.nodes[c.id as usize].neighbors[0] {
                    if visited.insert(neighbor) {
                        let d = self.quantized_dist(neighbor, query);
                        let furthest_dist = results.peek().map(|f| f.0.dist).unwrap_or(f32::MAX);
                        if d < furthest_dist || results.len() < ef {
                            candidates.push(Candidate { id: neighbor, dist: d });
                            if filter.is_none() || filter.unwrap().contains(neighbor) {
                                results.push(FurthestCandidate(Candidate { id: neighbor, dist: d }));
                                if results.len() > ef {
                                    results.pop();
                                }
                            }
                        }
                    }
                }
            }
        }

        let mut result: Vec<Candidate> = results.into_iter().map(|f| f.0).collect();
        result.sort_by(|a, b| a.dist.partial_cmp(&b.dist).unwrap_or(Ordering::Equal));
        result
    }

    fn brute_force_search(&self, query: &[f32], k: usize, filter: &roaring::RoaringBitmap) -> Vec<(u32, f32)> {
        let mut results: Vec<(u32, f32)> = filter.iter()
            .filter(|&id| (id as usize) < self.vectors.len())
            .map(|id| (id, self.quantized_dist(id, query)))
            .collect();
        results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
        results.truncate(k);
        results
    }

    /// Distance between stored vector `a` and query. Uses quantized
    /// distance when available, falling back to float32.
    #[inline]
    fn quantized_dist(&self, a: u32, query: &[f32]) -> f32 {
        if let Some(ref qv) = self.quantized {
            qv.asymmetric_distance(a as usize, query, self.cached_query_norm.get())
        } else {
            super::distance(&self.vectors[a as usize], query, self.params.metric)
        }
    }

    /// Exact float32 distance (for reranking after quantized search).
    fn exact_dist(&self, a: u32, query: &[f32]) -> f32 {
        super::distance(&self.vectors[a as usize], query, self.params.metric)
    }

    pub fn len(&self) -> usize { self.vectors.len() }
    pub fn is_empty(&self) -> bool { self.vectors.is_empty() }
    pub fn dims(&self) -> usize { self.params.dims }

    /// Get a vector by ID.
    pub fn vector(&self, id: u32) -> &[f32] {
        &self.vectors[id as usize]
    }

    /// Serialize to bytes.
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut buf = Vec::new();
        // Header
        buf.extend_from_slice(&(self.params.dims as u32).to_le_bytes());
        buf.extend_from_slice(&(self.params.m as u32).to_le_bytes());
        buf.push(self.params.metric as u8);
        buf.extend_from_slice(&(self.vectors.len() as u32).to_le_bytes());
        buf.extend_from_slice(&self.entry_point.unwrap_or(u32::MAX).to_le_bytes());
        buf.extend_from_slice(&(self.max_level as u32).to_le_bytes());

        // Vectors: flat f32 array
        for v in &self.vectors {
            for &f in v {
                buf.extend_from_slice(&f.to_le_bytes());
            }
        }

        // Nodes: level + per-level neighbor lists
        for node in &self.nodes {
            buf.extend_from_slice(&(node.level as u32).to_le_bytes());
            buf.extend_from_slice(&(node.neighbors.len() as u32).to_le_bytes());
            for layer in &node.neighbors {
                buf.extend_from_slice(&(layer.len() as u32).to_le_bytes());
                for &n in layer {
                    buf.extend_from_slice(&n.to_le_bytes());
                }
            }
        }

        // Quantized vectors (if present)
        if let Some(ref qv) = self.quantized {
            buf.push(1u8); // has quantized
            let qbytes = qv.to_bytes();
            buf.extend_from_slice(&(qbytes.len() as u32).to_le_bytes());
            buf.extend_from_slice(&qbytes);
        } else {
            buf.push(0u8); // no quantized
        }

        buf
    }

    /// Deserialize from bytes.
    pub fn from_bytes(data: &[u8]) -> Self {
        let mut pos = 0;
        let dims = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize; pos += 4;
        let m = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize; pos += 4;
        let metric_byte = data[pos]; pos += 1;
        let metric = match metric_byte {
            0 => DistanceMetric::Cosine,
            1 => DistanceMetric::DotProduct,
            _ => DistanceMetric::L2,
        };
        let num_vectors = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize; pos += 4;
        let ep = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()); pos += 4;
        let entry_point = if ep == u32::MAX { None } else { Some(ep) };
        let max_level = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize; pos += 4;

        // Vectors
        let mut vectors = Vec::with_capacity(num_vectors);
        for _ in 0..num_vectors {
            let mut v = Vec::with_capacity(dims);
            for _ in 0..dims {
                v.push(f32::from_le_bytes(data[pos..pos+4].try_into().unwrap()));
                pos += 4;
            }
            vectors.push(v);
        }

        // Nodes
        let mut nodes = Vec::with_capacity(num_vectors);
        for _ in 0..num_vectors {
            let level = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize; pos += 4;
            let num_layers = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize; pos += 4;
            let mut neighbors = Vec::with_capacity(num_layers);
            for _ in 0..num_layers {
                let num_neighbors = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize; pos += 4;
                let mut layer = Vec::with_capacity(num_neighbors);
                for _ in 0..num_neighbors {
                    layer.push(u32::from_le_bytes(data[pos..pos+4].try_into().unwrap())); pos += 4;
                }
                neighbors.push(layer);
            }
            nodes.push(Node { neighbors, level });
        }

        // Quantized vectors (optional, backward compatible)
        let quantized = if pos < data.len() && data[pos] == 1 {
            pos += 1;
            let qlen = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize; pos += 4;
            let qv = super::quantize::QuantizedVectors::from_bytes(&data[pos..pos+qlen]);
            Some(qv)
        } else {
            // No quantized data or old format — quantize from float32 vectors
            if !vectors.is_empty() {
                Some(super::quantize::QuantizedVectors::quantize(&vectors, metric))
            } else {
                None
            }
        };

        Self {
            params: HnswParams { dims, m, ef_construction: 100, metric },
            vectors, nodes, entry_point, max_level,
            quantized,
            cached_query_norm: std::cell::Cell::new(0.0),
        }
    }
}

// --- Zero-copy HNSW searcher for scalable segment access ---

/// Parsed graph structure (nodes + neighbors). This is the only part
/// that needs heap allocation. Cached in SegmentReader.
pub struct ParsedGraph {
    pub params: HnswParams,
    pub num_vectors: usize,
    pub entry_point: Option<u32>,
    pub max_level: usize,
    pub nodes: Vec<Node>,
    /// Byte offset where vector data starts in the segment slice.
    pub vector_data_offset: usize,
    /// Byte offset where quantized data starts (if present).
    pub quantized_offset: Option<usize>,
    /// Length of quantized data block.
    pub quantized_len: usize,
}

impl ParsedGraph {
    /// Parse the graph structure from HNSW bytes. Records byte offsets
    /// for vector and quantized data but does NOT copy them.
    pub fn parse(data: &[u8]) -> Self {
        let mut pos = 0;
        let dims = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize; pos += 4;
        let m = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize; pos += 4;
        let metric_byte = data[pos]; pos += 1;
        let metric = match metric_byte {
            0 => DistanceMetric::Cosine,
            1 => DistanceMetric::DotProduct,
            _ => DistanceMetric::L2,
        };
        let num_vectors = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize; pos += 4;
        let ep = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()); pos += 4;
        let entry_point = if ep == u32::MAX { None } else { Some(ep) };
        let max_level = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize; pos += 4;

        // Record vector data offset, skip past vector bytes
        let vector_data_offset = pos;
        pos += num_vectors * dims * 4; // skip float32 vectors

        // Parse nodes (graph structure — this we DO parse into memory)
        let mut nodes = Vec::with_capacity(num_vectors);
        for _ in 0..num_vectors {
            let level = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize; pos += 4;
            let num_layers = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize; pos += 4;
            let mut neighbors = Vec::with_capacity(num_layers);
            for _ in 0..num_layers {
                let num_neighbors = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize; pos += 4;
                let mut layer = Vec::with_capacity(num_neighbors);
                for _ in 0..num_neighbors {
                    layer.push(u32::from_le_bytes(data[pos..pos+4].try_into().unwrap())); pos += 4;
                }
                neighbors.push(layer);
            }
            nodes.push(Node { neighbors, level });
        }

        // Record quantized data offset
        let (quantized_offset, quantized_len) = if pos < data.len() && data[pos] == 1 {
            pos += 1;
            let qlen = u32::from_le_bytes(data[pos..pos+4].try_into().unwrap()) as usize;
            pos += 4;
            (Some(pos), qlen)
        } else {
            (None, 0)
        };

        Self {
            params: HnswParams { dims, m, ef_construction: 100, metric },
            num_vectors, entry_point, max_level, nodes,
            vector_data_offset, quantized_offset, quantized_len,
        }
    }
}

/// Zero-copy HNSW searcher that borrows segment data for vectors and
/// quantized data. Only the graph structure (via `ParsedGraph`) is
/// heap-allocated.
pub struct HnswSearcher<'a> {
    graph: &'a ParsedGraph,
    /// Raw HNSW bytes (vectors + nodes + quantized, borrowed from segment).
    data: &'a [u8],
    /// Parsed quantized calibration (mins, scales, norms) — small, cached.
    quantized_cal: Option<QuantizedCalibration<'a>>,
    cached_query_norm: std::cell::Cell<f32>,
}

/// Quantized calibration data (borrowed from segment bytes).
struct QuantizedCalibration<'a> {
    dims: usize,
    num_vectors: usize,
    metric: DistanceMetric,
    mins: &'a [u8],    // dims × 4 bytes (f32 le)
    scales: &'a [u8],  // dims × 4 bytes (f32 le)
    norms: &'a [u8],   // num_vectors × 4 bytes (f32 le)
    data: &'a [u8],    // num_vectors × dims bytes (u8 quantized values)
}

impl<'a> HnswSearcher<'a> {
    /// Create a searcher borrowing segment data and a parsed graph.
    pub fn new(data: &'a [u8], graph: &'a ParsedGraph) -> Self {
        let quantized_cal = graph.quantized_offset.map(|qoff| {
            let qdata = &data[qoff..qoff + graph.quantized_len];
            // Parse quantized calibration structure (borrow slices)
            let dims = graph.params.dims;
            let num_vectors = graph.num_vectors;
            let mut p = 0;
            // Skip dims, num_vectors, metric (already in graph.params)
            p += 4 + 4 + 1; // dims u32 + num_vectors u32 + metric u8
            let mins = &qdata[p..p + dims * 4]; p += dims * 4;
            let scales = &qdata[p..p + dims * 4]; p += dims * 4;
            let norms = &qdata[p..p + num_vectors * 4]; p += num_vectors * 4;
            let vdata = &qdata[p..p + num_vectors * dims];
            QuantizedCalibration { dims, num_vectors, metric: graph.params.metric, mins, scales, norms, data: vdata }
        });

        Self {
            graph, data, quantized_cal,
            cached_query_norm: std::cell::Cell::new(0.0),
        }
    }

    /// Search for k nearest neighbors.
    pub fn search(&self, query: &[f32], k: usize, ef: usize) -> Vec<(u32, f32)> {
        self.search_filtered(query, k, ef, None)
    }

    pub fn search_filtered(
        &self, query: &[f32], k: usize, ef: usize,
        filter: Option<&roaring::RoaringBitmap>,
    ) -> Vec<(u32, f32)> {
        if self.graph.entry_point.is_none() {
            return Vec::new();
        }

        if self.graph.params.metric == DistanceMetric::Cosine {
            self.cached_query_norm.set(
                query.iter().map(|x| x * x).sum::<f32>().sqrt()
            );
        }

        if let Some(bm) = filter {
            if (bm.len() as f64) < (self.graph.num_vectors as f64 * 0.01) {
                return self.brute_force_search(query, k, bm);
            }
        }

        let ep = self.graph.entry_point.unwrap();
        let ef_actual = ef.max(k);

        let mut current = ep;
        for lev in (1..=self.graph.max_level).rev() {
            current = self.greedy_closest(current, query, lev);
        }

        let candidates = self.search_layer_0(query, current, ef_actual, filter);

        // Rerank with exact float32 distance
        let mut results: Vec<(u32, f32)> = candidates.into_iter()
            .map(|c| (c.id, self.exact_dist(c.id, query)))
            .collect();
        results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
        results.truncate(k);
        results
    }

    /// Read a float32 value from the vector data region.
    #[inline]
    fn read_f32(&self, byte_offset: usize) -> f32 {
        f32::from_le_bytes(self.data[byte_offset..byte_offset+4].try_into().unwrap())
    }

    /// Exact float32 distance (for reranking).
    fn exact_dist(&self, idx: u32, query: &[f32]) -> f32 {
        let dims = self.graph.params.dims;
        let base = self.graph.vector_data_offset + (idx as usize) * dims * 4;
        let mut vec = Vec::with_capacity(dims);
        for d in 0..dims {
            vec.push(self.read_f32(base + d * 4));
        }
        super::distance(&vec, query, self.graph.params.metric)
    }

    /// Fast approximate distance (quantized if available, else float32).
    #[inline]
    fn approx_dist(&self, idx: u32, query: &[f32]) -> f32 {
        if let Some(ref cal) = self.quantized_cal {
            self.quantized_cosine_dist(cal, idx as usize, query)
        } else {
            self.exact_dist(idx, query)
        }
    }

    fn quantized_cosine_dist(&self, cal: &QuantizedCalibration, idx: usize, query: &[f32]) -> f32 {
        let dims = cal.dims;
        let qvec = &cal.data[idx * dims..(idx + 1) * dims];
        let mut dot = 0.0f32;
        for d in 0..dims {
            let min = f32::from_le_bytes(cal.mins[d*4..d*4+4].try_into().unwrap());
            let scale = f32::from_le_bytes(cal.scales[d*4..d*4+4].try_into().unwrap());
            let dequant = min + qvec[d] as f32 * scale;
            dot += dequant * query[d];
        }
        let norm_offset = idx * 4;
        let stored_norm = f32::from_le_bytes(cal.norms[norm_offset..norm_offset+4].try_into().unwrap());
        let query_norm = self.cached_query_norm.get();
        let denom = stored_norm * query_norm;
        if denom == 0.0 { 1.0 } else { 1.0 - dot / denom }
    }

    fn greedy_closest(&self, start: u32, query: &[f32], level: usize) -> u32 {
        let mut current = start;
        let mut current_dist = self.approx_dist(current, query);
        loop {
            let mut changed = false;
            if level < self.graph.nodes[current as usize].neighbors.len() {
                for &neighbor in &self.graph.nodes[current as usize].neighbors[level] {
                    let d = self.approx_dist(neighbor, query);
                    if d < current_dist {
                        current = neighbor;
                        current_dist = d;
                        changed = true;
                    }
                }
            }
            if !changed { break; }
        }
        current
    }

    fn search_layer_0(
        &self, query: &[f32], entry: u32, ef: usize,
        filter: Option<&roaring::RoaringBitmap>,
    ) -> Vec<Candidate> {
        let mut visited = HashSet::new();
        let mut candidates = BinaryHeap::new();
        let mut results = BinaryHeap::new();

        let d = self.approx_dist(entry, query);
        visited.insert(entry);
        candidates.push(Candidate { id: entry, dist: d });
        if filter.is_none() || filter.unwrap().contains(entry) {
            results.push(FurthestCandidate(Candidate { id: entry, dist: d }));
        }

        while let Some(c) = candidates.pop() {
            let furthest_dist = results.peek().map(|f| f.0.dist).unwrap_or(f32::MAX);
            if c.dist > furthest_dist && results.len() >= ef { break; }

            if !self.graph.nodes[c.id as usize].neighbors.is_empty() {
                for &neighbor in &self.graph.nodes[c.id as usize].neighbors[0] {
                    if visited.insert(neighbor) {
                        let d = self.approx_dist(neighbor, query);
                        let furthest_dist = results.peek().map(|f| f.0.dist).unwrap_or(f32::MAX);
                        if d < furthest_dist || results.len() < ef {
                            candidates.push(Candidate { id: neighbor, dist: d });
                            if filter.is_none() || filter.unwrap().contains(neighbor) {
                                results.push(FurthestCandidate(Candidate { id: neighbor, dist: d }));
                                if results.len() > ef { results.pop(); }
                            }
                        }
                    }
                }
            }
        }

        let mut result: Vec<Candidate> = results.into_iter().map(|f| f.0).collect();
        result.sort_by(|a, b| a.dist.partial_cmp(&b.dist).unwrap_or(Ordering::Equal));
        result
    }

    fn brute_force_search(&self, query: &[f32], k: usize, filter: &roaring::RoaringBitmap) -> Vec<(u32, f32)> {
        let mut results: Vec<(u32, f32)> = filter.iter()
            .filter(|&id| (id as usize) < self.graph.num_vectors)
            .map(|id| (id, self.exact_dist(id, query)))
            .collect();
        results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
        results.truncate(k);
        results
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_params(dims: usize) -> HnswParams {
        HnswParams { dims, m: 8, ef_construction: 50, metric: DistanceMetric::L2 }
    }

    #[test]
    fn build_and_search_small() {
        let mut builder = HnswBuilder::new(make_params(2));
        // 4 points in a square
        builder.add(vec![0.0, 0.0]); // 0
        builder.add(vec![1.0, 0.0]); // 1
        builder.add(vec![0.0, 1.0]); // 2
        builder.add(vec![1.0, 1.0]); // 3

        let index = builder.build();
        let results = index.search(&[0.1, 0.1], 2, 10);
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].0, 0); // closest to (0,0)
    }

    #[test]
    fn search_returns_k_results() {
        let mut builder = HnswBuilder::new(make_params(3));
        for i in 0..20 {
            builder.add(vec![i as f32, 0.0, 0.0]);
        }
        let index = builder.build();
        let results = index.search(&[5.0, 0.0, 0.0], 5, 20);
        assert_eq!(results.len(), 5);
        // Should include point 5 (exact match)
        assert!(results.iter().any(|(id, _)| *id == 5));
    }

    #[test]
    fn recall_on_random_data() {
        let dims = 16;
        let n = 200;
        let mut builder = HnswBuilder::new(HnswParams {
            dims, m: 12, ef_construction: 50, metric: DistanceMetric::L2,
        });

        // Deterministic "random" vectors
        let mut rng: u64 = 12345;
        let mut vectors = Vec::new();
        for _ in 0..n {
            let mut v = Vec::with_capacity(dims);
            for _ in 0..dims {
                rng ^= rng << 13; rng ^= rng >> 7; rng ^= rng << 17;
                v.push((rng as f32 / u64::MAX as f32) * 2.0 - 1.0);
            }
            vectors.push(v);
        }

        for v in &vectors {
            builder.add(v.clone());
        }
        let index = builder.build();

        // Query with the first vector — it should find itself
        let results = index.search(&vectors[0], 1, 50);
        assert_eq!(results[0].0, 0);

        // Compute recall@10 against brute-force
        let query = &vectors[42];
        let hnsw_results = index.search(query, 10, 50);

        // Brute-force top-10
        let mut brute: Vec<(u32, f32)> = (0..n as u32)
            .map(|i| (i, super::super::distance(&vectors[i as usize], query, DistanceMetric::L2)))
            .collect();
        brute.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
        let brute_top10: HashSet<u32> = brute[..10].iter().map(|x| x.0).collect();
        let hnsw_top10: HashSet<u32> = hnsw_results.iter().map(|x| x.0).collect();

        let recall = brute_top10.intersection(&hnsw_top10).count() as f64 / 10.0;
        assert!(recall >= 0.8, "recall@10 = {recall}, expected >= 0.8");
    }

    #[test]
    fn filtered_search() {
        let mut builder = HnswBuilder::new(make_params(2));
        for i in 0..10 {
            builder.add(vec![i as f32, 0.0]);
        }
        let index = builder.build();

        // Only allow even IDs
        let mut filter = roaring::RoaringBitmap::new();
        for i in (0..10).step_by(2) {
            filter.insert(i);
        }

        let results = index.search_filtered(&[3.0, 0.0], 3, 20, Some(&filter));
        // All results should have even IDs
        for (id, _) in &results {
            assert!(id % 2 == 0, "filtered result should be even, got {id}");
        }
        // Closest even to 3.0 is 2 or 4
        assert!(results[0].0 == 2 || results[0].0 == 4);
    }

    #[test]
    fn serialization_round_trip() {
        let mut builder = HnswBuilder::new(make_params(3));
        builder.add(vec![1.0, 2.0, 3.0]);
        builder.add(vec![4.0, 5.0, 6.0]);
        builder.add(vec![7.0, 8.0, 9.0]);
        let index = builder.build();

        let bytes = index.to_bytes();
        let restored = HnswIndex::from_bytes(&bytes);

        assert_eq!(restored.len(), 3);
        assert_eq!(restored.dims(), 3);

        let r1 = index.search(&[1.0, 2.0, 3.0], 1, 10);
        let r2 = restored.search(&[1.0, 2.0, 3.0], 1, 10);
        assert_eq!(r1[0].0, r2[0].0);
    }

    #[test]
    fn empty_index() {
        let builder = HnswBuilder::new(make_params(2));
        let index = builder.build();
        let results = index.search(&[0.0, 0.0], 5, 10);
        assert!(results.is_empty());
    }

    #[test]
    fn single_vector() {
        let mut builder = HnswBuilder::new(make_params(2));
        builder.add(vec![1.0, 1.0]);
        let index = builder.build();
        let results = index.search(&[0.0, 0.0], 1, 10);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0, 0);
    }

    #[test]
    fn cosine_metric() {
        let params = HnswParams { dims: 3, m: 8, ef_construction: 50, metric: DistanceMetric::Cosine };
        let mut builder = HnswBuilder::new(params);
        builder.add(vec![1.0, 0.0, 0.0]); // 0: x-axis
        builder.add(vec![0.0, 1.0, 0.0]); // 1: y-axis
        builder.add(vec![0.9, 0.1, 0.0]); // 2: near x-axis

        let index = builder.build();
        let results = index.search(&[1.0, 0.0, 0.0], 2, 10);
        // Closest by cosine to x-axis: 0 (identical) then 2 (near x-axis)
        assert_eq!(results[0].0, 0);
        assert_eq!(results[1].0, 2);
    }
}
