/// HNSW (Hierarchical Navigable Small World) — graph-based approximate nearest-neighbour search.
///
/// Each vector is a node. During insertion, it is connected to `M` neighbours at its entry layer
/// and below. Search starts from a fixed entry point at the top layer and greedily descends,
/// maintaining a candidate heap of size `ef` at each layer.
///
/// Vectors are stored quantized (packed bytes) here, same as BruteForce.
/// The graph structure (layer assignments + adjacency lists) lives in `nodes`.

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

// ---------------------------------------------------------------------------
// Node
// ---------------------------------------------------------------------------

struct Node {
    id:          u64,
    packed:      Vec<u8>,
    norm:        f32,
    /// Adjacency lists per layer. Each entry: (neighbor_idx, connection_score).
    /// Score is stored for SELECT-NEIGHBORS pruning during construction only.
    neighbours:  Vec<Vec<(u32, f32)>>,
}

// ---------------------------------------------------------------------------
// HNSW index
// ---------------------------------------------------------------------------

pub struct Hnsw {
    nodes:         Vec<Node>,
    entry:         Option<u32>, // index into nodes for current entry point (top layer)
    entry_layer:   usize,
    m:             usize,       // max neighbours per layer (base layer uses 2*m)
    ef_construction: usize,
    pub ef_search: usize,       // runtime-tunable
    pub packed_bytes: usize,
    ml:            f64,         // level multiplier = 1/ln(M)
    // Stored during add() phase for exact FP32 neighbor selection.
    // Freed by finalize_build() — never serialized.
    build_vecs:    Vec<Vec<f32>>,
}

impl Hnsw {
    pub fn new(m: usize, ef_construction: usize, ef_search: usize, packed_bytes: usize) -> Self {
        let ml = 1.0 / (m as f64).ln().max(1e-9);
        Self {
            nodes: Vec::new(),
            entry: None,
            entry_layer: 0,
            m,
            ef_construction,
            ef_search,
            packed_bytes,
            ml,
            build_vecs: Vec::new(),
        }
    }

    pub fn len(&self) -> usize { self.nodes.len() }

    /// Insert a new vector into the graph.
    ///
    /// `transformed`: RHDH-transformed f32 vector used for distance computation.
    /// `metric`: used to score candidates during graph construction (must match search metric).
    /// `score_fn(a_packed, b_packed) -> f32`: asymmetric similarity (higher = closer).
    pub fn insert<F>(
        &mut self,
        id:          u64,
        packed:      Vec<u8>,
        norm:        f32,
        transformed: &[f32],
        metric:      Metric,
        _score_fn:   &F,
    ) where F: Fn(&[f32], &[u8]) -> f32 {
        let level     = self.random_level();
        let new_idx   = self.nodes.len() as u32;

        // Pre-allocate adjacency lists for each layer.
        let max_neighbours: Vec<Vec<(u32, f32)>> = (0..=level)
            .map(|l| Vec::with_capacity(self.m_at(l)))
            .collect();

        self.nodes.push(Node { id, packed, norm, neighbours: max_neighbours });
        // Keep FP32 copy for exact neighbor scoring during graph construction.
        self.build_vecs.push(transformed.to_vec());

        if self.entry.is_none() {
            self.entry       = Some(new_idx);
            self.entry_layer = level;
            return;
        }

        let ep = self.entry.unwrap();

        // Phase 1: top layers above `level` — greedy descent, ef=1.
        let mut ep_idx = ep;
        let top = self.entry_layer;
        for lc in (level + 1..=top).rev() {
            let candidates = Self::search_layer_build(
                &self.nodes, &self.build_vecs, transformed, ep_idx, 1, lc, metric,
            );
            if let Some(best) = candidates.into_iter().next() {
                ep_idx = best.idx;
            }
        }

        // Phase 2: layers 0..=level — build connections.
        let ef_c = self.ef_construction;
        for lc in (0..=level.min(top)).rev() {
            let m_lc = self.m_at(lc);
            let candidates = Self::search_layer_build(
                &self.nodes, &self.build_vecs, transformed, ep_idx, ef_c, lc, metric,
            );

            // Select M best candidates (with scores for SELECT-NEIGHBORS pruning).
            let selected: Vec<(u32, f32)> = candidates.iter()
                .take(m_lc)
                .map(|c| (c.idx, c.score))
                .collect();

            self.nodes[new_idx as usize].neighbours[lc] = selected.clone();

            for &(nb, score) in &selected {
                let m_max   = self.m_at(lc);
                let nb_nbrs = &mut self.nodes[nb as usize].neighbours[lc];
                if !nb_nbrs.iter().any(|&(idx, _)| idx == new_idx) {
                    nb_nbrs.push((new_idx, score));
                    // SELECT-NEIGHBORS: if degree exceeds M_max, keep the M_max
                    // highest-scored connections (Malkov & Yashunin 2018, Alg. 1).
                    if nb_nbrs.len() > m_max {
                        nb_nbrs.sort_unstable_by(|a, b|
                            b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal));
                        nb_nbrs.truncate(m_max);
                    }
                }
            }

            if let Some(best) = candidates.into_iter().next() {
                ep_idx = best.idx;
            }
        }

        // Update entry point if new node is on a higher layer.
        if level > self.entry_layer {
            self.entry       = Some(new_idx);
            self.entry_layer = level;
        }
    }

    /// Search for top-K nearest neighbours.
    ///
    /// `transformed_query`: RHDH-transformed query.
    pub fn search<F>(
        &self,
        transformed_query: &[f32],
        k:                 usize,
        metric:            Metric,
        score_fn:          &F,
    ) -> Vec<SearchResult>
    where F: Fn(&[f32], &[u8]) -> f32 {
        if self.nodes.is_empty() || k == 0 { return Vec::new(); }

        let ep   = self.entry.unwrap();
        let mut ep_idx = ep;

        // Greedy descent through top layers.
        for lc in (1..=self.entry_layer).rev() {
            let candidates = self.search_layer(transformed_query, ep_idx, 1, lc, metric, score_fn);
            if let Some(best) = candidates.into_iter().next() {
                ep_idx = best.idx;
            }
        }

        // Base layer search with ef_search candidates.
        let ef = self.ef_search.max(k);
        let candidates = self.search_layer(transformed_query, ep_idx, ef, 0, metric, score_fn);

        // Scores already adjusted during traversal — use directly.
        let mut results: Vec<SearchResult> = candidates.into_iter().take(k).map(|c| {
            let node = &self.nodes[c.idx as usize];
            SearchResult { id: node.id, score: c.score }
        }).collect();
        results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(Ordering::Equal));
        results
    }

    // ---------------------------------------------------------------------------
    // Internal helpers
    // ---------------------------------------------------------------------------

    fn search_layer<F>(
        &self,
        query:    &[f32],
        ep_idx:   u32,
        ef:       usize,
        layer:    usize,
        metric:   Metric,
        score_fn: &F,
    ) -> Vec<ScoredNode>
    where F: Fn(&[f32], &[u8]) -> f32 {
        let ep_score = self.raw_score_adjusted(query, ep_idx as usize, metric, score_fn);

        // Candidate max-heap (best first); W min-heap (worst first for eviction).
        let mut candidates: BinaryHeap<ScoredNode> = BinaryHeap::new();
        let mut visited:    HashSet<u32>            = HashSet::new();
        let mut result:     Vec<ScoredNode>         = Vec::with_capacity(ef);

        candidates.push(ScoredNode { score: ep_score, idx: ep_idx });
        visited.insert(ep_idx);
        result.push(ScoredNode { score: ep_score, idx: ep_idx });

        let mut explored: usize = 0;
        while let Some(current) = candidates.pop() {
            explored += 1;
            let worst_result = result.iter().map(|r| r.score).fold(f32::INFINITY, f32::min);
            // Only terminate after exploring at least ef candidates — guarantees ef controls
            // search depth even when quantization noise compresses score gaps.
            if current.score < worst_result && explored > ef {
                break;
            }

            for &(nb, _) in self.neighbours_at(current.idx as usize, layer) {
                if visited.contains(&nb) { continue; }
                visited.insert(nb);

                let nb_score = self.raw_score_adjusted(query, nb as usize, metric, score_fn);

                // Always add to candidates so exploration reaches ef depth.
                candidates.push(ScoredNode { score: nb_score, idx: nb });
                // Only add to result if competitive.
                if result.len() < ef || nb_score > worst_result {
                    result.push(ScoredNode { score: nb_score, idx: nb });
                    if result.len() > ef {
                        result.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(Ordering::Equal));
                        result.truncate(ef);
                    }
                }
            }
        }

        result.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(Ordering::Equal));
        result
    }

    #[inline]
    fn raw_score_adjusted<F>(&self, query: &[f32], node_idx: usize, metric: Metric, score_fn: &F) -> f32
    where F: Fn(&[f32], &[u8]) -> f32 {
        let node = &self.nodes[node_idx];
        let raw  = score_fn(query, &node.packed);
        metric.adjust_score(raw, node.norm)
    }

    /// Release FP32 build vectors after graph construction is complete.
    /// Frees ~(n × dim × 4) bytes that are no longer needed after build().
    pub fn finalize_build(&mut self) {
        self.build_vecs.clear();
        self.build_vecs.shrink_to_fit();
    }

    /// Exact FP32 search over already-inserted nodes, metric-aware.
    /// Used only during insert() — `nodes` and `build_vecs` are borrowed separately
    /// so that `insert()` can mutate `nodes` after the call returns.
    fn search_layer_build(
        nodes:      &[Node],
        build_vecs: &[Vec<f32>],
        query:      &[f32],
        ep_idx:     u32,
        ef:         usize,
        layer:      usize,
        metric:     Metric,
    ) -> Vec<ScoredNode> {
        let ep_score = fp32_build_score(query, &build_vecs[ep_idx as usize], metric);

        let mut candidates: BinaryHeap<ScoredNode> = BinaryHeap::new();
        let mut visited:    HashSet<u32>            = HashSet::new();
        let mut result:     Vec<ScoredNode>         = Vec::with_capacity(ef);

        candidates.push(ScoredNode { score: ep_score, idx: ep_idx });
        visited.insert(ep_idx);
        result.push(ScoredNode { score: ep_score, idx: ep_idx });

        let mut explored: usize = 0;
        while let Some(current) = candidates.pop() {
            explored += 1;
            let worst_result = result.iter().map(|r| r.score).fold(f32::INFINITY, f32::min);
            if current.score < worst_result && explored > ef {
                break;
            }

            let nbs: &[(u32, f32)] = {
                let n = &nodes[current.idx as usize];
                if layer < n.neighbours.len() { &n.neighbours[layer] } else { &[] }
            };
            for &(nb, _) in nbs {
                if visited.contains(&nb) { continue; }
                visited.insert(nb);

                let nb_score = fp32_build_score(query, &build_vecs[nb as usize], metric);
                candidates.push(ScoredNode { score: nb_score, idx: nb });
                if result.len() < ef || nb_score > worst_result {
                    result.push(ScoredNode { score: nb_score, idx: nb });
                    if result.len() > ef {
                        result.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(Ordering::Equal));
                        result.truncate(ef);
                    }
                }
            }
        }

        result.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(Ordering::Equal));
        result
    }

    #[inline]
    fn neighbours_at(&self, node_idx: usize, layer: usize) -> &[(u32, f32)] {
        let n = &self.nodes[node_idx];
        if layer < n.neighbours.len() { &n.neighbours[layer] } else { &[] }
    }

    /// Max neighbours allowed at a given layer (base layer = 2*M, upper layers = M).
    #[inline]
    fn m_at(&self, layer: usize) -> usize {
        if layer == 0 { self.m * 2 } else { self.m }
    }

    /// Random layer for a new node: geometric distribution with mean `1/ln(M)`.
    /// Deterministic based on node count (no external RNG needed — sequential inserts).
    fn random_level(&self) -> usize {
        let n = self.nodes.len() as u64 + 1;
        // LCG-derived uniform in [0,1) using node count as seed.
        let r = lcg_uniform(n);
        (-r.ln() * self.ml) as usize
    }

    // ---------------------------------------------------------------------------
    // Serialization
    // ---------------------------------------------------------------------------

    /// Write HNSW graph data after the main header section.
    ///
    /// Layout (v5+):
    ///   [NODE_COUNT: u32]
    ///   [ENTRY:      u32] — node index of entry point (u32::MAX if empty)
    ///   [ENTRY_LYR:  u32] — entry layer
    ///   For each node:
    ///     [ID:        u64]
    ///     [PACKED:    u8 × pb]
    ///     [NORM:      f32]
    ///     [N_LAYERS:  u8]
    ///     For each layer:
    ///       [N_NB:    u16]  — neighbour count (u8 in v4, changed to u16 in v5)
    ///       [NBS:     u32 × N_NB]
    pub fn write_data<W: std::io::Write>(&self, w: &mut W) -> std::io::Result<()> {
        let n = self.nodes.len() as u32;
        w.write_all(&n.to_le_bytes())?;
        let entry = self.entry.unwrap_or(u32::MAX);
        w.write_all(&entry.to_le_bytes())?;
        w.write_all(&(self.entry_layer as u32).to_le_bytes())?;
        for node in &self.nodes {
            w.write_all(&node.id.to_le_bytes())?;
            w.write_all(&node.packed)?;
            w.write_all(&node.norm.to_le_bytes())?;
            let n_layers = node.neighbours.len() as u8;
            w.write_all(&[n_layers])?;
            for layer_nbs in &node.neighbours {
                let n_nb = layer_nbs.len() as u16;
                w.write_all(&n_nb.to_le_bytes())?;
                for &(nb, _) in layer_nbs {
                    w.write_all(&nb.to_le_bytes())?;
                }
            }
        }
        Ok(())
    }

    /// Deserialize HNSW graph written by `write_data`.
    pub fn read_data<R: std::io::Read>(
        r:           &mut R,
        packed_bytes: usize,
        m:           usize,
        ef_construction: usize,
        ef_search:   usize,
    ) -> std::io::Result<Self> {
        let mut b4 = [0u8; 4];
        let mut b8 = [0u8; 8];

        r.read_exact(&mut b4)?;
        let n = u32::from_le_bytes(b4) as usize;
        r.read_exact(&mut b4)?;
        let entry_raw = u32::from_le_bytes(b4);
        r.read_exact(&mut b4)?;
        let entry_layer = u32::from_le_bytes(b4) as usize;

        let mut nodes = Vec::with_capacity(n);
        for _ in 0..n {
            r.read_exact(&mut b8)?;
            let id = u64::from_le_bytes(b8);
            let mut packed = vec![0u8; packed_bytes];
            r.read_exact(&mut packed)?;
            r.read_exact(&mut b4)?;
            let norm = f32::from_le_bytes(b4);
            let mut n_layers_buf = [0u8; 1];
            r.read_exact(&mut n_layers_buf)?;
            let n_layers = n_layers_buf[0] as usize;
            let mut neighbours = Vec::with_capacity(n_layers);
            for _ in 0..n_layers {
                let mut n_nb_buf = [0u8; 2];
                r.read_exact(&mut n_nb_buf)?;
                let n_nb = u16::from_le_bytes(n_nb_buf) as usize;
                let mut nbs = Vec::with_capacity(n_nb);
                for _ in 0..n_nb {
                    r.read_exact(&mut b4)?;
                    nbs.push((u32::from_le_bytes(b4), 0.0f32));
                }
                neighbours.push(nbs);
            }
            nodes.push(Node { id, packed, norm, neighbours });
        }

        let entry = if entry_raw == u32::MAX { None } else { Some(entry_raw) };
        let ml    = 1.0 / (m as f64).ln().max(1e-9);

        Ok(Self { nodes, entry, entry_layer, m, ef_construction, ef_search, packed_bytes, ml,
                  build_vecs: Vec::new() })
    }
}

// ---------------------------------------------------------------------------
// FP32 exact dot product for build-time graph construction.
// ---------------------------------------------------------------------------

#[inline]
fn fp32_dot(a: &[f32], b: &[f32]) -> f32 {
    a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
}

/// FP32 build-time score matching the search metric.
///
/// Cosine / Dot: plain dot product (cosine vectors are unit-norm after transform).
/// L2: `dot(q, v) - 0.5 * ||v||²` — equivalent to `-(1/2)||q-v||²` + constant.
///     Using the wrong metric (plain dot) during build causes the graph topology to
///     mismatch the search metric, degrading recall on unnormalized vectors.
#[inline]
fn fp32_build_score(q: &[f32], v: &[f32], metric: Metric) -> f32 {
    let d = fp32_dot(q, v);
    match metric {
        Metric::L2 => d - 0.5 * fp32_dot(v, v),
        _          => d,
    }
}

// ---------------------------------------------------------------------------
// LCG-based deterministic "random" level — no rand dependency needed.
// ---------------------------------------------------------------------------

fn lcg_uniform(seed: u64) -> f64 {
    // Knuth's LCG constants.
    let x = seed.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
    // Map high 53 bits to [0, 1).
    (x >> 11) as f64 / (1u64 << 53) as f64
}

// ---------------------------------------------------------------------------
// Heap entry
// ---------------------------------------------------------------------------

#[derive(PartialEq)]
struct ScoredNode { score: f32, idx: u32 }
impl Eq for ScoredNode {}
impl PartialOrd for ScoredNode {
    fn partial_cmp(&self, o: &Self) -> Option<Ordering> { Some(self.cmp(o)) }
}
impl Ord for ScoredNode {
    fn cmp(&self, o: &Self) -> Ordering {
        self.score.partial_cmp(&o.score)
            .unwrap_or(Ordering::Equal)
            .then(self.idx.cmp(&o.idx))
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    fn unit_vec(dim: usize, seed: u64) -> Vec<f32> {
        use rand::SeedableRng;
        use rand::RngCore;
        use rand_chacha::ChaCha20Rng;
        let mut rng = ChaCha20Rng::seed_from_u64(seed);
        let mut v: Vec<f32> = (0..dim)
            .map(|_| (rng.next_u32() as f32 / u32::MAX as f32) * 2.0 - 1.0)
            .collect();
        let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        v.iter_mut().for_each(|x| *x /= norm);
        v
    }

    fn dot(a: &[f32], b: &[u8]) -> f32 {
        // Trivial score: interpret packed bytes as f32 scaled indices — just for tests.
        // The test uses transformed vecs as "packed" for simplicity.
        let _ = b;
        let _ = a;
        0.0
    }

    #[test]
    fn insert_and_len() {
        let mut hnsw = Hnsw::new(8, 50, 20, 4);
        for i in 0..50u64 {
            hnsw.insert(i, vec![0u8; 4], 1.0, &unit_vec(16, i), Metric::Cosine, &|_q, _p| 0.0f32);
        }
        assert_eq!(hnsw.len(), 50);
    }

    #[test]
    fn graph_has_neighbours() {
        let dim = 32usize;
        let mut hnsw = Hnsw::new(8, 50, 20, dim);
        let vecs: Vec<Vec<f32>> = (0..100u64).map(|i| unit_vec(dim, i)).collect();
        for (i, v) in vecs.iter().enumerate() {
            // Use the f32 vec both as "packed" (fake) and transformed.
            let packed = vec![i as u8; dim];
            hnsw.insert(i as u64, packed, 1.0, v, Metric::Cosine, &|_q, _p| 0.0f32);
        }
        // Most nodes in base layer should have at least 1 neighbour.
        let nodes_with_nbs = hnsw.nodes.iter()
            .filter(|n| !n.neighbours.is_empty() && !n.neighbours[0].is_empty())
            .count();
        assert!(nodes_with_nbs > 50, "expected most nodes to have neighbours, got {}", nodes_with_nbs);
    }

    #[test]
    fn search_returns_k() {
        let dim = 32usize;
        let mut hnsw = Hnsw::new(8, 50, 20, dim);
        let vecs: Vec<Vec<f32>> = (0..100u64).map(|i| unit_vec(dim, i)).collect();
        for (i, v) in vecs.iter().enumerate() {
            let packed = vec![0u8; dim];
            hnsw.insert(i as u64, packed, 1.0, v, Metric::Cosine, &|_q, _p| 0.0f32);
        }
        let query = unit_vec(dim, 999);
        let results = hnsw.search(&query, 10, Metric::Cosine, &|_q, _p| 0.0f32);
        assert_eq!(results.len(), 10);
    }

    #[test]
    fn max_neighbours_bounded() {
        // With M=4, base-layer M_max = 8. Insert enough nodes to trigger
        // back-connection accumulation and verify the invariant holds.
        let dim = 16usize;
        let m   = 4usize;
        let mut hnsw = Hnsw::new(m, 50, 20, dim);
        let vecs: Vec<Vec<f32>> = (0..200u64).map(|i| unit_vec(dim, i)).collect();
        for (i, v) in vecs.iter().enumerate() {
            hnsw.insert(i as u64, vec![0u8; dim], 1.0, v, Metric::Cosine, &|_q, _p| 0.0f32);
        }
        let m_max_base = m * 2;
        for (node_idx, node) in hnsw.nodes.iter().enumerate() {
            if !node.neighbours.is_empty() {
                assert!(
                    node.neighbours[0].len() <= m_max_base,
                    "node {} has {} neighbours in layer 0 (max {})",
                    node_idx, node.neighbours[0].len(), m_max_base
                );
            }
            for (lc, layer) in node.neighbours.iter().enumerate().skip(1) {
                assert!(
                    layer.len() <= m,
                    "node {} has {} neighbours in layer {} (max {})",
                    node_idx, layer.len(), lc, m
                );
            }
        }
    }

    #[test]
    fn save_load_roundtrip() {
        let dim = 16usize;
        let mut hnsw = Hnsw::new(4, 20, 10, dim);
        for i in 0..20u64 {
            hnsw.insert(i, vec![i as u8; dim], 1.0, &unit_vec(dim, i), Metric::Cosine, &|_q, _p| 0.0f32);
        }
        let mut buf = Vec::new();
        hnsw.write_data(&mut buf).unwrap();

        let loaded = Hnsw::read_data(
            &mut std::io::Cursor::new(&buf), dim, 4, 20, 10
        ).unwrap();
        assert_eq!(loaded.len(), 20);
        assert_eq!(loaded.entry, hnsw.entry);
        assert_eq!(loaded.entry_layer, hnsw.entry_layer);
    }
}
