//! HNSW insertion operations
//!
//! Implements single insert, batch insert, and neighbor selection heuristic.

use super::HNSWIndex;
use crate::vector::hnsw::error::{HNSWError, Result};
use ordered_float::OrderedFloat;
use tracing::{debug, error, info, instrument};

impl HNSWIndex {
    /// Validate vector dimensions and values for insertion
    fn validate_insert_vector(&self, vector: &[f32]) -> Result<()> {
        if vector.len() != self.dimensions() {
            error!(
                expected_dim = self.dimensions(),
                actual_dim = vector.len(),
                "Dimension mismatch during insert"
            );
            return Err(HNSWError::DimensionMismatch {
                expected: self.dimensions(),
                actual: vector.len(),
            });
        }

        if vector.iter().any(|x| !x.is_finite()) {
            error!("Invalid vector: contains NaN or Inf values");
            return Err(HNSWError::InvalidVector);
        }

        Ok(())
    }

    /// Store vector and create node, returns (node_id, level)
    fn store_and_create_node(&mut self, vector: &[f32]) -> Result<(u32, u8)> {
        // Allocate node in unified storage
        let node_id = self.storage.allocate_node();

        // Assign random level
        let level = self.random_level();

        // Store vector in colocated storage
        self.storage
            .set_vector(node_id, vector)
            .map_err(HNSWError::storage)?;

        // Set node metadata
        self.storage.set_slot(node_id, node_id); // slot == node_id for new nodes
        self.storage.set_level(node_id, level);

        // Allocate upper level storage if needed
        if level > 0 {
            self.storage.allocate_upper_levels(node_id, level);
        }

        Ok((node_id, level))
    }

    /// Update entry point if new node has higher level
    fn maybe_update_entry_point(&mut self, node_id: u32, level: u8) -> Result<()> {
        let entry_point_id = self
            .entry_point
            .ok_or_else(|| HNSWError::internal("Entry point should exist after first insert"))?;
        let entry_level = self.storage.level(entry_point_id);

        if level > entry_level {
            self.entry_point = Some(node_id);
            debug!(
                old_entry = entry_point_id,
                new_entry = node_id,
                old_level = entry_level,
                new_level = level,
                "Updated entry point to higher level node"
            );
        }

        Ok(())
    }

    /// Insert a vector into the index
    ///
    /// Returns the node ID assigned to this vector.
    #[instrument(skip(self, vector), fields(dimensions = vector.len(), index_size = self.len()))]
    pub fn insert(&mut self, vector: &[f32]) -> Result<u32> {
        self.validate_insert_vector(vector)?;
        let (node_id, level) = self.store_and_create_node(vector)?;

        // If this is the first node, set as entry point and return
        if self.entry_point.is_none() {
            self.entry_point = Some(node_id);
            return Ok(node_id);
        }

        self.insert_into_graph(node_id, vector, level)?;
        self.maybe_update_entry_point(node_id, level)?;

        debug!(
            node_id = node_id,
            level = level,
            index_size = self.len(),
            "Successfully inserted vector"
        );

        Ok(node_id)
    }

    /// Add bidirectional links and prune neighbors if necessary
    fn reconcile_node_neighbors(
        &mut self,
        node_id: u32,
        neighbors: &[u32],
        level: u8,
    ) -> Result<()> {
        let m = self.params.m_for_level(level);

        // Set neighbors for this node
        self.storage
            .set_neighbors_at_level(node_id, level, neighbors.to_vec());

        // Add reverse links to neighbors, pruning if necessary
        // The key insight is that we must add the new node and prune properly,
        // not skip adding when at capacity (which breaks the graph).
        for &neighbor_id in neighbors {
            // Use zero-copy COW for the contains check: avoids cloning when the
            // neighbor is already connected (the common continue path).
            let existing = self.storage.neighbors_at_level_cow(neighbor_id, level);

            // Skip if already connected (avoid duplicates)
            if existing.contains(&node_id) {
                continue;
            }

            // Need to mutate: materialize the owned Vec only when adding a link.
            let mut neighbor_neighbors = existing.into_owned();

            // Add the new reverse link
            neighbor_neighbors.push(node_id);

            // Prune if over capacity
            if neighbor_neighbors.len() > m {
                let neighbor_vec = self
                    .storage
                    .get_dequantized(neighbor_id)
                    .ok_or(HNSWError::VectorNotFound(neighbor_id))?;
                neighbor_neighbors = self.select_neighbors_heuristic(
                    neighbor_id,
                    &neighbor_neighbors,
                    m,
                    level,
                    &neighbor_vec,
                )?;
            }

            // Update the neighbor's connections
            self.storage
                .set_neighbors_at_level(neighbor_id, level, neighbor_neighbors);
        }

        Ok(())
    }

    /// Insert a vector with entry point hints for faster insertion
    ///
    /// Used by graph merging to speed up insertion when we already know
    /// good starting points (neighbors from the source graph).
    ///
    /// # Arguments
    /// * `vector` - Vector to insert
    /// * `entry_hints` - Node IDs to use as starting points (must exist in index)
    /// * `ef` - Expansion factor for search (lower = faster, may reduce quality)
    ///
    /// # Performance
    /// ~5x faster than standard insert when hints are good neighbors
    #[instrument(skip(self, vector, entry_hints), fields(dimensions = vector.len(), hints = entry_hints.len()))]
    pub fn insert_with_hints(
        &mut self,
        vector: &[f32],
        entry_hints: &[u32],
        ef: usize,
    ) -> Result<u32> {
        self.validate_insert_vector(vector)?;
        let (node_id, level) = self.store_and_create_node(vector)?;

        // If this is the first node, set as entry point and return
        if self.entry_point.is_none() {
            self.entry_point = Some(node_id);
            return Ok(node_id);
        }

        // Filter hints to valid node IDs that exist in the index
        let valid_hints: Vec<u32> = entry_hints
            .iter()
            .filter(|&&id| (id as usize) < self.storage.len())
            .copied()
            .collect();

        // If no valid hints, fall back to standard insertion
        if valid_hints.is_empty() {
            return self
                .insert_into_graph(node_id, vector, level)
                .map(|()| node_id);
        }

        // Use hints as starting points for graph insertion
        self.insert_into_graph_with_hints(node_id, vector, level, &valid_hints, ef)?;
        self.maybe_update_entry_point(node_id, level)?;

        Ok(node_id)
    }

    /// Insert node into graph using entry hints instead of global entry point
    pub(super) fn insert_into_graph_with_hints(
        &mut self,
        node_id: u32,
        vector: &[f32],
        level: u8,
        entry_hints: &[u32],
        ef: usize,
    ) -> Result<()> {
        // Start search from hints (skip upper layer traversal)
        let mut nearest = entry_hints.to_vec();

        // Insert at levels 0..=level (iterate from top to bottom)
        // Use full precision distances during graph construction for better quality
        for lc in (0..=level).rev() {
            // Find ef nearest neighbors at this level using reduced ef
            // Use distance-aware search to avoid recomputing in heuristic
            let candidates_with_distances =
                self.search_layer_for_construction(vector, &nearest, ef, lc)?;

            // Select M best neighbors using heuristic (distances already computed)
            let m = self.params.m_for_level(lc);
            let neighbors =
                self.select_neighbors_heuristic_with_distances(&candidates_with_distances, m)?;

            self.reconcile_node_neighbors(node_id, &neighbors, lc)?;

            // Update nearest for next level
            nearest = candidates_with_distances
                .iter()
                .map(|(id, _)| *id)
                .collect();
        }

        Ok(())
    }

    /// Batch insert multiple vectors with batched sequential graph construction
    ///
    /// This method reduces overhead versus repeated single inserts by:
    /// 1. Storing all vectors first (no graph construction)
    /// 2. Building the HNSW graph in one batched pass
    ///
    /// # Performance
    /// - Small batches (<100): Use `insert()` for simplicity
    /// - Medium/large batches amortize allocation and bookkeeping costs
    /// - For true parallel construction, use `HNSWIndex::build_parallel()`
    ///
    /// # Algorithm
    /// - Pre-allocate all nodes and levels (deterministic)
    /// - Build graph connections in a single batched pass
    /// - Reuse stored vectors and inline reverse-link pruning
    ///
    /// # Arguments
    /// * `vectors` - Batch of vectors to insert
    ///
    /// # Returns
    /// Vector of node IDs corresponding to inserted vectors
    #[instrument(skip(self, vectors), fields(batch_size = vectors.len()))]
    pub fn batch_insert(&mut self, vectors: Vec<Vec<f32>>) -> Result<Vec<u32>> {
        use rayon::prelude::*;
        use std::sync::atomic::{AtomicU32, Ordering};

        if vectors.is_empty() {
            return Ok(Vec::new());
        }

        let batch_size = vectors.len();
        info!(batch_size, "Starting batched HNSW insertion");

        // Parallel validation (fast, no graph modifications)
        let dimensions = self.dimensions();
        let validation_start = std::time::Instant::now();

        vectors.par_iter().try_for_each(|vec| -> Result<()> {
            if vec.len() != dimensions {
                return Err(HNSWError::DimensionMismatch {
                    expected: dimensions,
                    actual: vec.len(),
                });
            }
            if vec.iter().any(|x| !x.is_finite()) {
                return Err(HNSWError::InvalidVector);
            }
            Ok(())
        })?;

        debug!(
            duration_ms = validation_start.elapsed().as_millis(),
            "Parallel validation complete"
        );

        // Phase 1: Store all vectors and create nodes (fast, sequential)
        let storage_start = std::time::Instant::now();
        let mut node_ids = Vec::with_capacity(batch_size);
        let mut node_levels = Vec::with_capacity(batch_size);

        // Track highest level node in this batch for entry point update AFTER graph construction
        let mut highest_level_node: Option<(u32, u8)> = None;

        for vector in &vectors {
            // Allocate node in unified storage
            let node_id = self.storage.allocate_node();

            // Assign level (deterministic from RNG state)
            let level = self.random_level();

            // Store vector
            self.storage
                .set_vector(node_id, vector)
                .map_err(HNSWError::storage)?;

            // Set node metadata
            self.storage.set_slot(node_id, node_id);
            self.storage.set_level(node_id, level);

            // Allocate upper level storage if needed
            if level > 0 {
                self.storage.allocate_upper_levels(node_id, level);
            }

            node_ids.push(node_id);
            node_levels.push(level);

            // Track highest level node (entry point update deferred until after graph construction)
            if self.entry_point.is_none() {
                // First node ever - set entry point immediately
                self.entry_point = Some(node_id);
                highest_level_node = Some((node_id, level));
            } else {
                // Track highest level for later update
                match highest_level_node {
                    None => highest_level_node = Some((node_id, level)),
                    Some((_, prev_level)) if level > prev_level => {
                        highest_level_node = Some((node_id, level));
                    }
                    _ => {}
                }
            }
        }

        debug!(
            duration_ms = storage_start.elapsed().as_millis(),
            nodes_added = node_ids.len(),
            "Vector storage complete"
        );

        // Phase 2: Build graph connections in a single batched pass
        let graph_start = std::time::Instant::now();

        // If this is the only node, no graph to build
        if self.storage.len() == 1 {
            info!("Single node, no graph construction needed");
            return Ok(node_ids);
        }

        // Batched graph construction
        let nodes_to_insert: Vec<(u32, u8)> = node_ids
            .iter()
            .zip(node_levels.iter())
            .map(|(&id, &level)| (id, level))
            .collect();

        // Use atomic counter for progress tracking
        let progress_counter = AtomicU32::new(0);
        let progress_interval = if batch_size >= 1000 {
            batch_size / 10
        } else {
            batch_size
        };

        // Sequential insertion into graph.
        // NodeStorage is Send+Sync safe, but this path mutates neighbor lists directly
        // and therefore does not use the lock-based ParallelBuilder implementation.
        for (node_id, level) in &nodes_to_insert {
            // Get vector for this node
            let vector = self
                .storage
                .get_dequantized(*node_id)
                .ok_or(HNSWError::VectorNotFound(*node_id))?;

            // Build graph connections
            let entry_point = self.entry_point.ok_or(HNSWError::EmptyIndex)?;
            let entry_level = self.storage.level(entry_point);

            // Search for nearest neighbors at each level above target level
            let mut nearest = vec![entry_point];
            for lc in ((*level + 1)..=entry_level).rev() {
                nearest = self
                    .search_layer_for_construction(&vector, &nearest, 1, lc)?
                    .into_iter()
                    .map(|(id, _)| id)
                    .collect();
            }

            // Insert at levels 0..=level
            for lc in (0..=*level).rev() {
                // Find ef_construction nearest neighbors at this level
                // Use distance-aware search to avoid recomputing in heuristic
                let candidates_with_distances = self.search_layer_for_construction(
                    &vector,
                    &nearest,
                    self.params.ef_construction,
                    lc,
                )?;

                // Select M best neighbors using heuristic (distances already computed)
                let m = self.params.m_for_level(lc);
                let neighbors =
                    self.select_neighbors_heuristic_with_distances(&candidates_with_distances, m)?;

                // Set neighbors for this node
                self.storage
                    .set_neighbors_at_level(*node_id, lc, neighbors.clone());

                // Add reverse links with proper pruning
                for &neighbor_id in &neighbors {
                    // Zero-copy COW: skip clone when already connected (common path).
                    let existing = self.storage.neighbors_at_level_cow(neighbor_id, lc);
                    if !existing.contains(node_id) {
                        let mut neighbor_neighbors = existing.into_owned();
                        neighbor_neighbors.push(*node_id);
                        // Prune if over capacity (must recompute distances for pruning)
                        if neighbor_neighbors.len() > m
                            && let Some(neighbor_vec) = self.storage.get_dequantized(neighbor_id)
                            && let Ok(pruned) = self.select_neighbors_heuristic(
                                neighbor_id,
                                &neighbor_neighbors,
                                m,
                                lc,
                                &neighbor_vec,
                            )
                        {
                            neighbor_neighbors = pruned;
                        }
                        self.storage
                            .set_neighbors_at_level(neighbor_id, lc, neighbor_neighbors);
                    }
                }

                // Update nearest for next level
                nearest = candidates_with_distances
                    .iter()
                    .map(|(id, _)| *id)
                    .collect();
            }

            // Progress tracking
            let count = progress_counter.fetch_add(1, Ordering::Relaxed) + 1;
            if count.is_multiple_of(progress_interval as u32) {
                let elapsed = graph_start.elapsed().as_secs_f64();
                let rate = count as f64 / elapsed;
                info!(
                    progress = count,
                    total = batch_size,
                    percent = (count as usize * 100) / batch_size,
                    rate_vec_per_sec = rate as u64,
                    "Batched graph construction progress"
                );
            }
        }

        // Update entry point AFTER graph construction (critical for incremental inserts)
        if let Some((new_entry, new_level)) = highest_level_node
            && let Some(current_entry) = self.entry_point
        {
            let current_level = self.storage.level(current_entry);
            if new_level > current_level {
                self.entry_point = Some(new_entry);
            }
        }

        // Note: Post-insert pruning pass removed - inline pruning during reverse link
        // addition (lines 412-424) already ensures neighbors never exceed M.
        // This eliminates redundant O(n) pass over all nodes.

        let total_time = graph_start.elapsed().as_secs_f64();
        let final_rate = batch_size as f64 / total_time;

        info!(
            inserted = node_ids.len(),
            duration_secs = total_time,
            rate_vec_per_sec = final_rate as u64,
            "Batched insertion complete"
        );

        Ok(node_ids)
    }

    /// Insert node into graph structure
    ///
    /// Implements HNSW insertion algorithm (Malkov & Yashunin 2018)
    pub(super) fn insert_into_graph(
        &mut self,
        node_id: u32,
        vector: &[f32],
        level: u8,
    ) -> Result<()> {
        let entry_point = self.entry_point.ok_or(HNSWError::EmptyIndex)?;
        let entry_level = self.storage.level(entry_point);

        // Search for nearest neighbors at each level above target level
        // Use full precision distances during graph construction for better quality
        let mut nearest = vec![entry_point];
        for lc in ((level + 1)..=entry_level).rev() {
            nearest = self
                .search_layer_for_construction(vector, &nearest, 1, lc)?
                .into_iter()
                .map(|(id, _)| id)
                .collect();
        }
        for lc in (0..=level).rev() {
            // Find ef_construction nearest neighbors at this level
            // Use distance-aware search to avoid recomputing in heuristic
            let candidates_with_distances = self.search_layer_for_construction(
                vector,
                &nearest,
                self.params.ef_construction,
                lc,
            )?;

            // Select M best neighbors using heuristic (distances already computed)
            let m = self.params.m_for_level(lc);
            let neighbors =
                self.select_neighbors_heuristic_with_distances(&candidates_with_distances, m)?;

            self.reconcile_node_neighbors(node_id, &neighbors, lc)?;

            // Update nearest for next level (just node IDs for entry points)
            nearest = candidates_with_distances
                .iter()
                .map(|(id, _)| *id)
                .collect();
        }

        Ok(())
    }

    /// Select neighbors using heuristic (diverse neighbors, better recall)
    ///
    /// Algorithm from Malkov 2018, Section 4
    ///
    /// This version accepts (node_id, distance) pairs to avoid recomputing distances
    /// that were already calculated during the search phase.
    pub(super) fn select_neighbors_heuristic_with_distances(
        &self,
        candidates_with_distances: &[(u32, f32)],
        m: usize,
    ) -> Result<Vec<u32>> {
        if candidates_with_distances.len() <= m {
            return Ok(candidates_with_distances
                .iter()
                .map(|(id, _)| *id)
                .collect());
        }

        // Already sorted by distance from search_layer_with_distances
        let mut result = Vec::with_capacity(m);
        let mut remaining = Vec::new();

        // Heuristic: Select diverse neighbors
        for &(candidate_id, candidate_dist) in candidates_with_distances {
            if result.len() >= m {
                remaining.push(candidate_id);
                continue;
            }

            // Check if candidate is closer to query than to existing neighbors
            let mut good = true;
            for &result_id in &result {
                let dist_to_result = self.distance_between_cmp(candidate_id, result_id)?;
                if dist_to_result < candidate_dist {
                    good = false;
                    break;
                }
            }

            if good {
                result.push(candidate_id);
            } else {
                remaining.push(candidate_id);
            }
        }

        // Fill remaining slots with closest candidates if needed
        for candidate_id in remaining {
            if result.len() >= m {
                break;
            }
            result.push(candidate_id);
        }

        Ok(result)
    }

    /// Select neighbors using heuristic (legacy version that computes distances)
    ///
    /// Used by reconcile_node_neighbors for pruning reverse links where we
    /// don't have pre-computed distances.
    pub(super) fn select_neighbors_heuristic(
        &self,
        _node_id: u32,
        candidates: &[u32],
        m: usize,
        _level: u8,
        query_vector: &[f32],
    ) -> Result<Vec<u32>> {
        if candidates.len() <= m {
            return Ok(candidates.to_vec());
        }
        // Sort candidates by distance to query
        let mut sorted_candidates: Vec<_> = candidates
            .iter()
            .map(|&id| {
                let dist = self.distance_cmp(query_vector, id)?;
                Ok((id, dist))
            })
            .collect::<Result<Vec<_>>>()?;
        sorted_candidates.sort_unstable_by_key(|c| OrderedFloat(c.1));

        // Delegate to the distance-aware version
        self.select_neighbors_heuristic_with_distances(&sorted_candidates, m)
    }

    /// Add a neighbor to a node if there's room (neighbor count < M*2)
    ///
    /// Used for adding boundary connections during cluster merging.
    /// This is a unidirectional link - call twice for bidirectional.
    ///
    /// # Arguments
    /// * `node_id` - Node to add neighbor to
    /// * `neighbor_id` - Neighbor to add
    ///
    /// # Returns
    /// Ok(()) on success, even if neighbor wasn't added (already at capacity)
    pub fn add_neighbor_if_room(&mut self, node_id: u32, neighbor_id: u32) -> Result<()> {
        let level = 0u8; // Only level 0 for boundary connections
        let m_max = self.params.m_for_level(level);

        let neighbors = self.storage.neighbors(node_id);
        if neighbors.len() < m_max && !neighbors.contains(&neighbor_id) {
            self.storage.add_neighbor(node_id, level, neighbor_id);
        }

        Ok(())
    }
}
