//! HNSW-based vector index implementation.

use crate::error::{Result, StorageError};
use hnsw_rs::prelude::*;
use parking_lot::RwLock;
use std::collections::HashMap;
use stemedb_core::types::Hash;
use tracing::{debug, instrument, warn};

use super::VectorIndex;

/// Internal ID type for HNSW (maps to assertion hashes via lookup table).
type InternalId = usize;

/// HNSW-backed vector index implementation.
///
/// Uses the Hierarchical Navigable Small World algorithm for efficient
/// approximate nearest-neighbor search. Provides O(log N) query complexity
/// with high recall rates.
///
/// # Thread Safety
///
/// The index is protected by a RwLock, allowing concurrent reads but
/// exclusive writes. This matches the expected access pattern where
/// queries are frequent and inserts are batched during ingestion.
pub struct HnswVectorIndex {
    /// The HNSW graph structure.
    hnsw: RwLock<Hnsw<'static, f32, DistL2>>,

    /// Maps internal HNSW IDs to assertion hashes.
    id_to_hash: RwLock<HashMap<InternalId, Hash>>,

    /// Maps assertion hashes to internal HNSW IDs (for deduplication).
    hash_to_id: RwLock<HashMap<Hash, InternalId>>,

    /// The vector dimension this index was created for.
    dimension: usize,

    /// Counter for generating internal IDs.
    next_id: RwLock<InternalId>,
}

impl HnswVectorIndex {
    /// Default number of neighbors to connect during construction.
    /// Higher values = better recall but slower inserts.
    const DEFAULT_MAX_NB_CONNECTION: usize = 16;

    /// Default width of search during construction.
    /// Higher values = better recall but slower inserts.
    pub(crate) const DEFAULT_EF_CONSTRUCTION: usize = 200;

    /// Create a new vector index for the given dimension.
    ///
    /// # Arguments
    /// * `dimension` - The dimension of vectors to be indexed (e.g., 128, 768, 1536)
    ///
    /// # Panics
    /// Panics if dimension is 0.
    pub fn new(dimension: usize) -> Self {
        assert!(dimension > 0, "Vector dimension must be positive");

        let hnsw = Hnsw::new(
            Self::DEFAULT_MAX_NB_CONNECTION,
            10000, // Initial capacity, will grow as needed
            16,    // Number of layers
            Self::DEFAULT_EF_CONSTRUCTION,
            DistL2 {},
        );

        Self {
            hnsw: RwLock::new(hnsw),
            id_to_hash: RwLock::new(HashMap::new()),
            hash_to_id: RwLock::new(HashMap::new()),
            dimension,
            next_id: RwLock::new(0),
        }
    }

    /// Create a new vector index with custom HNSW parameters.
    ///
    /// # Arguments
    /// * `dimension` - The dimension of vectors to be indexed
    /// * `max_nb_connection` - Maximum number of neighbors per node (16-64 typical)
    /// * `ef_construction` - Search width during construction (200-800 typical)
    pub fn with_params(dimension: usize, max_nb_connection: usize, ef_construction: usize) -> Self {
        assert!(dimension > 0, "Vector dimension must be positive");
        assert!(max_nb_connection > 0, "max_nb_connection must be positive");
        assert!(ef_construction > 0, "ef_construction must be positive");

        let hnsw = Hnsw::new(max_nb_connection, 10000, 16, ef_construction, DistL2 {});

        Self {
            hnsw: RwLock::new(hnsw),
            id_to_hash: RwLock::new(HashMap::new()),
            hash_to_id: RwLock::new(HashMap::new()),
            dimension,
            next_id: RwLock::new(0),
        }
    }

    /// Check if an assertion hash is already in the index.
    pub fn contains(&self, hash: &Hash) -> bool {
        self.hash_to_id.read().contains_key(hash)
    }
}

impl VectorIndex for HnswVectorIndex {
    #[instrument(skip(self, vector), fields(hash = %hex::encode(hash), dim = vector.len()))]
    fn insert(&self, hash: &Hash, vector: &[f32]) -> Result<()> {
        // Validate dimension
        if vector.len() != self.dimension {
            return Err(StorageError::InputValidation(format!(
                "Vector dimension mismatch: expected {}, got {}",
                self.dimension,
                vector.len()
            )));
        }

        // Check for NaN/Inf values
        for (i, &v) in vector.iter().enumerate() {
            if v.is_nan() {
                return Err(StorageError::InputValidation(format!(
                    "Vector contains NaN at index {}",
                    i
                )));
            }
            if v.is_infinite() {
                return Err(StorageError::InputValidation(format!(
                    "Vector contains infinite value at index {}",
                    i
                )));
            }
        }

        // Check if already indexed (idempotency)
        {
            let hash_to_id = self.hash_to_id.read();
            if hash_to_id.contains_key(hash) {
                debug!(hash = %hex::encode(hash), "Vector already indexed, skipping");
                return Ok(());
            }
        }

        // Generate internal ID
        let internal_id = {
            let mut next_id = self.next_id.write();
            let id = *next_id;
            *next_id += 1;
            id
        };

        // Insert into HNSW
        {
            let hnsw = self.hnsw.write();
            // hnsw_rs expects a tuple (&[T], usize) for single insert
            hnsw.insert((vector, internal_id));
        }

        // Update lookup tables
        {
            let mut id_to_hash = self.id_to_hash.write();
            let mut hash_to_id = self.hash_to_id.write();
            id_to_hash.insert(internal_id, *hash);
            hash_to_id.insert(*hash, internal_id);
        }

        debug!(
            hash = %hex::encode(hash),
            internal_id,
            "Inserted vector into HNSW index"
        );

        Ok(())
    }

    #[instrument(skip(self, query), fields(dim = query.len(), k = k))]
    fn search(&self, query: &[f32], k: usize) -> Result<Vec<(Hash, f32)>> {
        // Validate dimension
        if query.len() != self.dimension {
            return Err(StorageError::InputValidation(format!(
                "Query dimension mismatch: expected {}, got {}",
                self.dimension,
                query.len()
            )));
        }

        // Check for NaN/Inf values in query
        for (i, &v) in query.iter().enumerate() {
            if v.is_nan() {
                return Err(StorageError::InputValidation(format!(
                    "Query vector contains NaN at index {}",
                    i
                )));
            }
            if v.is_infinite() {
                return Err(StorageError::InputValidation(format!(
                    "Query vector contains infinite value at index {}",
                    i
                )));
            }
        }

        // Handle empty index
        if self.is_empty() {
            debug!("Vector index is empty, returning no results");
            return Ok(Vec::new());
        }

        // Handle k=0
        if k == 0 {
            return Ok(Vec::new());
        }

        // Search HNSW
        let neighbors = {
            let hnsw = self.hnsw.read();
            // ef_search should be >= k for good recall
            let ef_search = k.max(Self::DEFAULT_EF_CONSTRUCTION / 2);
            hnsw.search(query, k, ef_search)
        };

        // Map internal IDs back to hashes
        let id_to_hash = self.id_to_hash.read();
        let mut results = Vec::with_capacity(neighbors.len());

        for neighbor in neighbors {
            let internal_id = neighbor.d_id;
            let distance = neighbor.distance;

            match id_to_hash.get(&internal_id) {
                Some(hash) => {
                    results.push((*hash, distance));
                }
                None => {
                    warn!(internal_id, "HNSW returned ID with no hash mapping, skipping");
                }
            }
        }

        debug!(results_count = results.len(), "Vector search complete");

        Ok(results)
    }

    fn dimension(&self) -> usize {
        self.dimension
    }

    fn len(&self) -> usize {
        self.id_to_hash.read().len()
    }
}
