// Copyright 2025 SochDB Authors
//
// Licensed under the Apache License, Version 2.0

//! Optimized HNSW search using CSR graph + internal IDs + batched expansion.
//!
//! This module provides a high-performance search implementation that combines:
//! - **Dense u32 IDs**: 4× less bandwidth, faster hashing
//! - **CSR adjacency**: Zero-alloc neighbor iteration
//! - **Bitmap visited set**: O(1) membership with bit ops
//! - **Batched expansion**: Process 4-8 candidates per iteration for locality
//! - **Software prefetch**: Hide DRAM latency with pipelined vector fetches
//!
//! # Performance Comparison
//!
//! | Operation | Original | Optimized | Improvement |
//! |-----------|----------|-----------|-------------|
//! | ID lookup | HashMap<u128> | Array[u32] | 3-4× |
//! | Visited check | HashSet insert | Bitmap bit | 5-10× |
//! | Neighbor iterate | SmallVec iter | Slice iter | 2-3× |
//! | Memory per node | ~600 bytes | ~80 bytes | 7× |
//!
//! # Usage
//!
//! ```rust,ignore
//! let search_view = index.create_optimized_search_view();
//! let results = search_view.search(&query, k, ef)?;
//! ```

use crate::csr_graph::{CsrGraph, InternalSearchCandidate};
use crate::internal_id::{IdMapper, InternalId, VisitedBitmap};
use std::cmp::Reverse;
use std::collections::BinaryHeap;

/// Batch size for frontier expansion.
/// 4-8 is optimal: balances locality gains vs extra work.
const BATCH_SIZE: usize = 4;

/// Cache line size for prefetch alignment.
#[allow(dead_code)]
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
const CACHE_LINE: usize = 64;

#[allow(dead_code)]
#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
const CACHE_LINE: usize = 64;

/// Optimized search view for HNSW graph.
///
/// This is an immutable snapshot optimized for search:
/// - Lock-free traversal
/// - Contiguous vector storage
/// - CSR adjacency
/// - Dense internal IDs
pub struct OptimizedSearchView {
    /// CSR graph structure.
    pub graph: CsrGraph,

    /// ID mapping (external ↔ internal).
    pub id_mapper: IdMapper,

    /// Vector storage: vectors[internal_id * dim .. (internal_id + 1) * dim].
    /// Stored as contiguous f32 for cache efficiency.
    pub vectors: Vec<f32>,

    /// Vector dimension.
    pub dimension: usize,

    /// Distance metric.
    pub metric: DistanceMetric,

    /// Maximum layer in the graph.
    pub max_layer: usize,

    /// Entry point for search.
    pub entry_point: Option<InternalId>,
}

/// Distance metric for vector similarity.
#[derive(Debug, Clone, Copy)]
pub enum DistanceMetric {
    Cosine,
    Euclidean,
    DotProduct,
}

impl OptimizedSearchView {
    /// Create a new optimized search view.
    pub fn new(
        graph: CsrGraph,
        id_mapper: IdMapper,
        vectors: Vec<f32>,
        dimension: usize,
        metric: DistanceMetric,
    ) -> Self {
        let entry_point = graph.entry_point;
        let max_layer = graph.max_layer;

        Self {
            graph,
            id_mapper,
            vectors,
            dimension,
            metric,
            max_layer,
            entry_point,
        }
    }

    /// Search for k nearest neighbors.
    ///
    /// # Arguments
    /// * `query` - Query vector
    /// * `k` - Number of neighbors to return
    /// * `ef` - Search expansion factor (higher = more accurate, slower)
    ///
    /// # Returns
    /// Vec of (external_id, distance) pairs, sorted by distance
    pub fn search(&self, query: &[f32], k: usize, ef: usize) -> Vec<(u128, f32)> {
        let entry_point = match self.entry_point {
            Some(ep) => ep,
            None => return Vec::new(),
        };

        // Search from top layer down to layer 1
        let mut curr_nearest = vec![InternalSearchCandidate {
            distance: self.calculate_distance(query, entry_point),
            id: entry_point,
        }];

        for layer in (1..=self.max_layer).rev() {
            curr_nearest = self.search_layer(query, &curr_nearest, 1, layer);
        }

        // Final search at layer 0 with full ef
        let candidates = self.search_layer_batched(query, &curr_nearest, ef.max(k), 0);

        // Convert to external IDs and return top k
        candidates
            .into_iter()
            .take(k)
            .filter_map(|c| {
                self.id_mapper
                    .to_external(c.id)
                    .map(|ext| (ext, c.distance))
            })
            .collect()
    }

    /// Single-candidate search layer (used for upper layers).
    fn search_layer(
        &self,
        query: &[f32],
        entry_points: &[InternalSearchCandidate],
        num_to_return: usize,
        layer: usize,
    ) -> Vec<InternalSearchCandidate> {
        let mut visited = VisitedBitmap::new(self.graph.num_nodes);
        let mut candidates = BinaryHeap::new();
        let mut results: BinaryHeap<Reverse<InternalSearchCandidate>> = BinaryHeap::new();

        // Initialize with entry points
        for ep in entry_points {
            visited.insert(ep.id);
            candidates.push(ep.clone());
            results.push(Reverse(ep.clone()));
            if results.len() > num_to_return {
                results.pop();
            }
        }

        while let Some(curr) = candidates.pop() {
            // Stop if current is farther than worst result
            if results.len() >= num_to_return {
                if let Some(Reverse(worst)) = results.peek() {
                    if curr.distance > worst.distance {
                        break;
                    }
                }
            }

            // Get neighbors from CSR (zero allocation)
            let neighbors = self.graph.neighbors(curr.id, layer);

            for &neighbor_id in neighbors {
                if visited.insert(neighbor_id) {
                    let distance = self.calculate_distance(query, neighbor_id);
                    let candidate = InternalSearchCandidate {
                        distance,
                        id: neighbor_id,
                    };

                    if results.len() < num_to_return {
                        candidates.push(candidate.clone());
                        results.push(Reverse(candidate));
                    } else if let Some(Reverse(worst)) = results.peek() {
                        if distance < worst.distance {
                            candidates.push(candidate.clone());
                            results.pop();
                            results.push(Reverse(candidate));
                        }
                    }
                }
            }
        }

        // Convert to sorted vec
        let mut sorted_results: Vec<_> = results.into_iter().map(|Reverse(c)| c).collect();
        sorted_results.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap());
        sorted_results
    }

    /// Batched frontier expansion for layer 0 (the hot path).
    ///
    /// Instead of processing one candidate at a time, we batch BATCH_SIZE candidates
    /// and process all their neighbors together. This improves:
    /// - Cache locality (neighbors are fetched together)
    /// - SIMD utilization (batch distance computations)
    /// - Prefetch effectiveness (more look-ahead)
    fn search_layer_batched(
        &self,
        query: &[f32],
        entry_points: &[InternalSearchCandidate],
        ef: usize,
        layer: usize,
    ) -> Vec<InternalSearchCandidate> {
        let mut visited = VisitedBitmap::new(self.graph.num_nodes);
        let mut candidates = BinaryHeap::new();
        let mut results: BinaryHeap<Reverse<InternalSearchCandidate>> = BinaryHeap::new();

        // Initialize with entry points
        for ep in entry_points {
            visited.insert(ep.id);
            candidates.push(ep.clone());
            results.push(Reverse(ep.clone()));
            if results.len() > ef {
                results.pop();
            }
        }

        // Scratch buffer for batch processing
        let mut batch: Vec<InternalSearchCandidate> = Vec::with_capacity(BATCH_SIZE);
        let mut neighbor_buffer: Vec<InternalId> = Vec::with_capacity(BATCH_SIZE * 32);

        while !candidates.is_empty() {
            // Get worst result distance for pruning
            let worst_distance = results
                .peek()
                .map(|Reverse(c)| c.distance)
                .unwrap_or(f32::MAX);

            // Collect batch of candidates
            batch.clear();
            while batch.len() < BATCH_SIZE {
                if let Some(curr) = candidates.pop() {
                    // Stop if current is farther than worst result
                    if results.len() >= ef && curr.distance > worst_distance {
                        // Put back for next iteration
                        candidates.push(curr);
                        break;
                    }
                    batch.push(curr);
                } else {
                    break;
                }
            }

            if batch.is_empty() {
                break;
            }

            // Collect all unvisited neighbors from batch
            neighbor_buffer.clear();
            for candidate in &batch {
                let neighbors = self.graph.neighbors(candidate.id, layer);

                // Prefetch vector data for neighbors we'll likely visit
                #[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
                {
                    for &neighbor_id in neighbors.iter().take(4) {
                        self.prefetch_vector(neighbor_id);
                    }
                }

                for &neighbor_id in neighbors {
                    if visited.insert(neighbor_id) {
                        neighbor_buffer.push(neighbor_id);
                    }
                }
            }

            // Process all neighbors (better locality than interleaved)
            let current_worst = results
                .peek()
                .map(|Reverse(c)| c.distance)
                .unwrap_or(f32::MAX);

            for &neighbor_id in &neighbor_buffer {
                let distance = self.calculate_distance(query, neighbor_id);

                if results.len() < ef {
                    let candidate = InternalSearchCandidate {
                        distance,
                        id: neighbor_id,
                    };
                    candidates.push(candidate.clone());
                    results.push(Reverse(candidate));
                } else if distance < current_worst {
                    let candidate = InternalSearchCandidate {
                        distance,
                        id: neighbor_id,
                    };
                    candidates.push(candidate.clone());
                    results.pop();
                    results.push(Reverse(candidate));
                }
            }
        }

        // Convert to sorted vec
        let mut sorted_results: Vec<_> = results.into_iter().map(|Reverse(c)| c).collect();
        sorted_results.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap());
        sorted_results
    }

    /// Calculate distance between query and stored vector.
    #[inline]
    fn calculate_distance(&self, query: &[f32], node: InternalId) -> f32 {
        let idx = node.get() as usize;
        let start = idx * self.dimension;
        let end = start + self.dimension;

        if end > self.vectors.len() {
            return f32::MAX;
        }

        let vector = &self.vectors[start..end];

        match self.metric {
            DistanceMetric::Euclidean => self.euclidean_distance(query, vector),
            DistanceMetric::Cosine => self.cosine_distance(query, vector),
            DistanceMetric::DotProduct => -self.dot_product(query, vector),
        }
    }

    /// Euclidean (L2) distance.
    #[inline]
    #[allow(unreachable_code)]
    fn euclidean_distance(&self, a: &[f32], b: &[f32]) -> f32 {
        // Use SIMD if available
        #[cfg(target_arch = "x86_64")]
        {
            if is_x86_feature_detected!("avx2") {
                return unsafe { self.euclidean_distance_avx2(a, b) };
            }
        }

        #[cfg(target_arch = "aarch64")]
        {
            return unsafe { self.euclidean_distance_neon(a, b) };
        }

        // Scalar fallback
        a.iter()
            .zip(b.iter())
            .map(|(x, y)| {
                let diff = x - y;
                diff * diff
            })
            .sum::<f32>()
            .sqrt()
    }

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

    /// Cosine distance (1 - cosine_similarity).
    #[inline]
    fn cosine_distance(&self, a: &[f32], b: &[f32]) -> f32 {
        let dot = self.dot_product(a, b);
        let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
        let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
        if norm_a == 0.0 || norm_b == 0.0 {
            return 1.0;
        }
        1.0 - (dot / (norm_a * norm_b))
    }

    /// Prefetch vector data for upcoming distance computation.
    #[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
    #[inline]
    fn prefetch_vector(&self, node: InternalId) {
        let idx = node.get() as usize;
        let start = idx * self.dimension;
        if start < self.vectors.len() {
            let ptr = self.vectors.as_ptr();
            // Prefetch first cache line of vector
            #[cfg(target_arch = "x86_64")]
            unsafe {
                use std::arch::x86_64::_mm_prefetch;
                _mm_prefetch(ptr.add(start) as *const i8, std::arch::x86_64::_MM_HINT_T0);
            }
            // For aarch64, use inline assembly for prefetch (stable)
            #[cfg(target_arch = "aarch64")]
            unsafe {
                let addr = ptr.add(start) as *const u8;
                std::arch::asm!(
                    "prfm pldl1keep, [{0}]",
                    in(reg) addr,
                    options(nostack, preserves_flags)
                );
            }
        }
    }

    /// AVX2-accelerated Euclidean distance.
    #[cfg(target_arch = "x86_64")]
    #[target_feature(enable = "avx2", enable = "fma")]
    unsafe fn euclidean_distance_avx2(&self, a: &[f32], b: &[f32]) -> f32 {
        use std::arch::x86_64::*;

        let len = a.len().min(b.len());
        let chunks = len / 8;
        let remainder = len % 8;

        let mut acc = _mm256_setzero_ps();

        for i in 0..chunks {
            let idx = i * 8;
            let va = _mm256_loadu_ps(a.as_ptr().add(idx));
            let vb = _mm256_loadu_ps(b.as_ptr().add(idx));
            let diff = _mm256_sub_ps(va, vb);
            acc = _mm256_fmadd_ps(diff, diff, acc);
        }

        // Horizontal sum
        let sum128 = _mm_add_ps(_mm256_castps256_ps128(acc), _mm256_extractf128_ps(acc, 1));
        let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128));
        let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1));
        let mut total = _mm_cvtss_f32(sum32);

        // Handle remainder
        for i in (chunks * 8)..(chunks * 8 + remainder) {
            let diff = a[i] - b[i];
            total += diff * diff;
        }

        total.sqrt()
    }

    /// NEON-accelerated Euclidean distance.
    #[cfg(target_arch = "aarch64")]
    #[target_feature(enable = "neon")]
    unsafe fn euclidean_distance_neon(&self, a: &[f32], b: &[f32]) -> f32 {
        use std::arch::aarch64::*;

        let len = a.len().min(b.len());
        let chunks = len / 4;
        let remainder = len % 4;

        // SAFETY: We're in an unsafe function with NEON enabled via target_feature
        let mut acc = vdupq_n_f32(0.0);

        for i in 0..chunks {
            let idx = i * 4;
            // SAFETY: NEON intrinsics are safe with valid pointers within bounds
            unsafe {
                let va = vld1q_f32(a.as_ptr().add(idx));
                let vb = vld1q_f32(b.as_ptr().add(idx));
                let diff = vsubq_f32(va, vb);
                acc = vfmaq_f32(acc, diff, diff);
            }
        }

        // Horizontal sum
        // SAFETY: vaddvq_f32 is safe with a valid NEON register
        let mut total = vaddvq_f32(acc);

        // Handle remainder
        for i in (chunks * 4)..(chunks * 4 + remainder) {
            let diff = a[i] - b[i];
            total += diff * diff;
        }

        total.sqrt()
    }

    /// Get number of vectors in the index.
    pub fn len(&self) -> usize {
        self.graph.num_nodes
    }

    /// Check if index is empty.
    pub fn is_empty(&self) -> bool {
        self.graph.num_nodes == 0
    }

    /// Get memory usage in bytes.
    pub fn memory_usage(&self) -> usize {
        let graph_mem = self.graph.memory_usage();
        let vector_mem = self.vectors.len() * 4;
        let mapper_mem = self.id_mapper.memory_usage();
        graph_mem + vector_mem + mapper_mem
    }
}

/// Builder for creating OptimizedSearchView from an existing index.
pub struct OptimizedSearchViewBuilder {
    dimension: usize,
    metric: DistanceMetric,
    max_degree: usize,
    max_degree_layer0: usize,
}

impl OptimizedSearchViewBuilder {
    /// Create a new builder.
    pub fn new(dimension: usize) -> Self {
        Self {
            dimension,
            metric: DistanceMetric::Euclidean,
            max_degree: 16,
            max_degree_layer0: 32,
        }
    }

    /// Set distance metric.
    pub fn metric(mut self, metric: DistanceMetric) -> Self {
        self.metric = metric;
        self
    }

    /// Set max connections per node.
    pub fn max_degree(mut self, m: usize, m0: usize) -> Self {
        self.max_degree = m;
        self.max_degree_layer0 = m0;
        self
    }

    /// Build from vectors and pre-computed edges.
    ///
    /// # Arguments
    /// * `external_ids` - External IDs for each vector
    /// * `vectors` - Flattened vector data (N × D)
    /// * `edges` - List of (layer, from_external_id, to_external_id) edges
    /// * `entry_point` - External ID of entry point
    pub fn build(
        self,
        external_ids: &[u128],
        vectors: &[f32],
        edges: &[(usize, u128, u128)],
        entry_point_external: u128,
    ) -> OptimizedSearchView {
        let id_mapper = IdMapper::with_capacity(external_ids.len());

        // Register all IDs
        for &ext_id in external_ids {
            id_mapper.register(ext_id);
        }

        // Find max layer
        let max_layer = edges.iter().map(|(l, _, _)| *l).max().unwrap_or(0);

        // Build CSR graph
        let mut builder = crate::csr_graph::CsrGraphBuilder::new(
            max_layer + 1,
            self.max_degree,
            self.max_degree_layer0,
        );

        for &(layer, from, to) in edges {
            if let (Some(from_int), Some(to_int)) =
                (id_mapper.to_internal(from), id_mapper.to_internal(to))
            {
                builder.add_edge(from_int, to_int, layer);
            }
        }

        if let Some(ep_int) = id_mapper.to_internal(entry_point_external) {
            builder.set_entry_point(ep_int);
        }

        let graph = builder.build();

        // Reorder vectors to match internal ID order
        let mut ordered_vectors = vec![0.0f32; external_ids.len() * self.dimension];
        for &ext_id in external_ids {
            if let Some(int_id) = id_mapper.to_internal(ext_id) {
                // Find original position (linear search, could optimize)
                if let Some(orig_pos) = external_ids.iter().position(|&id| id == ext_id) {
                    let src_start = orig_pos * self.dimension;
                    let dst_start = int_id.get() as usize * self.dimension;
                    ordered_vectors[dst_start..dst_start + self.dimension]
                        .copy_from_slice(&vectors[src_start..src_start + self.dimension]);
                }
            }
        }

        OptimizedSearchView::new(
            graph,
            id_mapper,
            ordered_vectors,
            self.dimension,
            self.metric,
        )
    }
}

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

    fn create_test_view() -> OptimizedSearchView {
        let dimension = 4;
        let num_vectors = 10;

        // Create simple test vectors
        let mut vectors = Vec::with_capacity(num_vectors * dimension);
        let mut external_ids = Vec::with_capacity(num_vectors);
        for i in 0..num_vectors {
            external_ids.push((i * 100) as u128);
            for d in 0..dimension {
                vectors.push((i * dimension + d) as f32 / 100.0);
            }
        }

        // Create edges: simple chain
        let mut edges = Vec::new();
        for i in 0..num_vectors - 1 {
            edges.push((0, external_ids[i], external_ids[i + 1]));
            edges.push((0, external_ids[i + 1], external_ids[i]));
        }

        OptimizedSearchViewBuilder::new(dimension)
            .metric(DistanceMetric::Euclidean)
            .build(&external_ids, &vectors, &edges, external_ids[0])
    }

    #[test]
    fn test_optimized_search_basic() {
        let view = create_test_view();
        assert_eq!(view.len(), 10);
        assert!(!view.is_empty());
    }

    #[test]
    fn test_optimized_search_query() {
        let view = create_test_view();

        // Query close to vector 0
        let query = vec![0.0, 0.01, 0.02, 0.03];
        let results = view.search(&query, 3, 10);

        assert!(!results.is_empty());
        // First result should be closest to vector 0 (external ID 0)
        assert_eq!(results[0].0, 0);
    }

    #[test]
    fn test_distance_calculations() {
        let view = create_test_view();

        let a = [1.0, 0.0, 0.0, 0.0];
        let b = [0.0, 1.0, 0.0, 0.0];

        let euclidean = view.euclidean_distance(&a, &b);
        assert!((euclidean - std::f32::consts::SQRT_2).abs() < 0.001);

        let dot = view.dot_product(&a, &b);
        assert!((dot - 0.0).abs() < 0.001);
    }

    #[test]
    fn test_memory_usage() {
        let view = create_test_view();
        let mem = view.memory_usage();

        // Should be reasonable for 10 vectors of dimension 4
        assert!(mem > 0);
        assert!(mem < 10000); // Less than 10KB
    }

    #[test]
    fn test_empty_search() {
        let graph = CsrGraph::new();
        let id_mapper = IdMapper::new();
        let view =
            OptimizedSearchView::new(graph, id_mapper, Vec::new(), 4, DistanceMetric::Euclidean);

        let query = vec![0.0, 0.0, 0.0, 0.0];
        let results = view.search(&query, 5, 10);
        assert!(results.is_empty());
    }

    #[test]
    fn test_batched_vs_single() {
        // Ensure batched search produces correct results
        let dimension = 8;
        let num_vectors = 100;

        let mut vectors = Vec::with_capacity(num_vectors * dimension);
        let mut external_ids = Vec::with_capacity(num_vectors);
        for i in 0..num_vectors {
            external_ids.push(i as u128);
            for d in 0..dimension {
                vectors.push(((i * dimension + d) as f32).sin());
            }
        }

        // Create edges: each node connects to 4 nearest by index
        let mut edges = Vec::new();
        for i in 0..num_vectors {
            for j in 1..=4 {
                if i + j < num_vectors {
                    edges.push((0, external_ids[i], external_ids[i + j]));
                    edges.push((0, external_ids[i + j], external_ids[i]));
                }
            }
        }

        let view = OptimizedSearchViewBuilder::new(dimension)
            .metric(DistanceMetric::Euclidean)
            .max_degree(16, 32)
            .build(&external_ids, &vectors, &edges, external_ids[50]);

        // Query should find reasonable results
        let query: Vec<f32> = (0..dimension)
            .map(|d| (50.0 * dimension as f32 + d as f32).sin())
            .collect();
        let results = view.search(&query, 10, 50);

        assert_eq!(results.len(), 10);
        // Results should be unique
        let unique: std::collections::HashSet<_> = results.iter().map(|(id, _)| id).collect();
        assert_eq!(unique.len(), 10);
    }
}
