//! Brute-force (exact) vector index.
//!
//! [`BruteForceIndex`] performs linear-scan L2 search over all stored vectors.
//! It is the correctness baseline: every other index implementation must return
//! the same top-k results (within quantization tolerance). It is also used for
//! small datasets where the O(n) scan cost is acceptable.

use std::collections::HashMap;
use std::io::Write;
use std::path::Path;
use std::sync::RwLock;

use super::{VectorError, VectorId, VectorIndex, VectorIndexConfig, VectorSearchResult};

// ---------------------------------------------------------------------------
// Binary format constants
// ---------------------------------------------------------------------------

/// Magic bytes identifying a brute-force vector index file.
const MAGIC: &[u8; 4] = b"BFVI";

/// Current binary format version.
const FORMAT_VERSION: u8 = 0x01;

// ---------------------------------------------------------------------------
// Distance functions
// ---------------------------------------------------------------------------

/// Compute the squared Euclidean (L2) distance between two vectors.
///
/// This avoids the `sqrt` call -- the squared distance preserves ranking order
/// and is sufficient for nearest-neighbor selection.
///
/// # Panics (debug only)
///
/// Debug-asserts that `a` and `b` have the same length.
pub fn l2_distance_sq(a: &[f32], b: &[f32]) -> f32 {
    debug_assert_eq!(a.len(), b.len());
    a.iter()
        .zip(b.iter())
        .map(|(x, y)| {
            let d = x - y;
            d * d
        })
        .sum()
}

// ---------------------------------------------------------------------------
// BruteForceIndex
// ---------------------------------------------------------------------------

/// Exact nearest-neighbor index using linear scan.
///
/// Stores all vectors in a `HashMap` behind a `RwLock`. Search computes L2
/// squared distance against every stored vector and returns the top-k by
/// ascending distance.
///
/// This implementation has no tombstones -- `delete()` is a true removal.
/// `len()` and `len_live()` always return the same value.
///
/// # Thread safety
///
/// All reads take a shared lock; all writes take an exclusive lock. This is
/// acceptable because brute-force is not a hot-path production index -- it is
/// used for correctness baselines, small datasets, and tests.
pub struct BruteForceIndex {
    vectors: RwLock<HashMap<VectorId, Vec<f32>>>,
    config: VectorIndexConfig,
}

impl BruteForceIndex {
    /// Create a new, empty brute-force index with the given configuration.
    #[must_use]
    pub fn new(config: VectorIndexConfig) -> Self {
        Self {
            vectors: RwLock::new(HashMap::new()),
            config,
        }
    }

    /// Validate that a vector's dimensionality matches the index configuration.
    const fn validate_dimensions(&self, vec: &[f32]) -> Result<(), VectorError> {
        if vec.len() != self.config.dimensions {
            return Err(VectorError::DimensionMismatch {
                expected: self.config.dimensions,
                got: vec.len(),
            });
        }
        Ok(())
    }
}

impl VectorIndex for BruteForceIndex {
    fn insert(&self, id: VectorId, embedding: &[f32]) -> Result<(), VectorError> {
        self.validate_dimensions(embedding)?;
        self.vectors
            .write()
            .map_err(|e| VectorError::Backend(format!("RwLock poisoned on write: {e}")))?
            .insert(id, embedding.to_vec());
        Ok(())
    }

    /// Search for the `k` nearest neighbors by exhaustive linear scan.
    ///
    /// The `ef_search` parameter is accepted for trait compliance but ignored --
    /// brute-force search is exact and has no beam width parameter.
    fn search(
        &self,
        query: &[f32],
        k: usize,
        _ef_search: usize,
    ) -> Result<Vec<VectorSearchResult>, VectorError> {
        self.validate_dimensions(query)?;
        let guard = self
            .vectors
            .read()
            .map_err(|e| VectorError::Backend(format!("RwLock poisoned on read: {e}")))?;

        let mut results: Vec<VectorSearchResult> = guard
            .iter()
            .map(|(id, vec)| VectorSearchResult {
                id: *id,
                distance: l2_distance_sq(query, vec),
            })
            .collect();

        drop(guard);
        results.sort_by(|a, b| {
            a.distance
                .partial_cmp(&b.distance)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        results.truncate(k);
        Ok(results)
    }

    /// Filtered search: only compute distance for vectors where `filter(id)` is true.
    ///
    /// The `ef_search` parameter is accepted for trait compliance but ignored.
    fn filtered_search(
        &self,
        query: &[f32],
        k: usize,
        _ef_search: usize,
        filter: &dyn Fn(VectorId) -> bool,
    ) -> Result<Vec<VectorSearchResult>, VectorError> {
        self.validate_dimensions(query)?;
        let guard = self
            .vectors
            .read()
            .map_err(|e| VectorError::Backend(format!("RwLock poisoned on read: {e}")))?;

        let mut results: Vec<VectorSearchResult> = guard
            .iter()
            .filter(|(id, _)| filter(**id))
            .map(|(id, vec)| VectorSearchResult {
                id: *id,
                distance: l2_distance_sq(query, vec),
            })
            .collect();

        drop(guard);
        results.sort_by(|a, b| {
            a.distance
                .partial_cmp(&b.distance)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        results.truncate(k);
        Ok(results)
    }

    fn delete(&self, id: VectorId) -> Result<(), VectorError> {
        let removed = self
            .vectors
            .write()
            .map_err(|e| VectorError::Backend(format!("RwLock poisoned on write: {e}")))?
            .remove(&id);
        if removed.is_none() {
            return Err(VectorError::NotFound { id });
        }
        Ok(())
    }

    /// Reserve additional capacity. This is a no-op for `HashMap`-backed storage --
    /// `HashMap` resizes automatically. The method is provided for trait compliance
    /// and always succeeds.
    fn reserve(&self, _additional: usize) -> Result<(), VectorError> {
        Ok(())
    }

    fn save(&self, path: &Path) -> Result<(), VectorError> {
        let guard = self
            .vectors
            .read()
            .map_err(|e| VectorError::Backend(format!("RwLock poisoned on read: {e}")))?;

        let mut file = std::fs::File::create(path)?;

        // Header: magic + version + dimensions + count
        file.write_all(MAGIC)?;
        file.write_all(&[FORMAT_VERSION])?;
        #[allow(clippy::cast_possible_truncation)]
        let dims = self.config.dimensions as u32;
        file.write_all(&dims.to_le_bytes())?;
        let count = guard.len() as u64;
        file.write_all(&count.to_le_bytes())?;

        // Per-vector: id + floats
        for (id, vec) in &*guard {
            file.write_all(&id.to_le_bytes())?;
            for &val in vec {
                file.write_all(&val.to_le_bytes())?;
            }
        }

        drop(guard);
        file.flush().map_err(std::convert::Into::into)
    }

    fn load(path: &Path, config: &VectorIndexConfig) -> Result<Self, VectorError> {
        let data = std::fs::read(path)?;
        Self::deserialize(&data, config)
    }

    /// For brute-force, `view` delegates to `load` -- there is no mmap mode.
    /// The entire index is read into memory regardless.
    fn view(path: &Path, config: &VectorIndexConfig) -> Result<Self, VectorError> {
        Self::load(path, config)
    }

    fn len(&self) -> usize {
        self.vectors.read().map_or(0, |guard| guard.len())
    }

    fn len_live(&self) -> usize {
        // No tombstones in brute-force -- all entries are live.
        self.len()
    }
}

impl BruteForceIndex {
    /// Deserialize from a byte buffer (shared by `load` and `view`).
    fn deserialize(data: &[u8], config: &VectorIndexConfig) -> Result<Self, VectorError> {
        // Minimum header size: 4 (magic) + 1 (version) + 4 (dims) + 8 (count) = 17
        const HEADER_SIZE: usize = 17;
        if data.len() < HEADER_SIZE {
            return Err(VectorError::CorruptedIndex(
                "file too small for header".into(),
            ));
        }

        // Validate magic
        if &data[..4] != MAGIC {
            return Err(VectorError::CorruptedIndex(format!(
                "invalid magic: expected {:?}, got {:?}",
                MAGIC,
                &data[..4]
            )));
        }

        // Validate version
        if data[4] != FORMAT_VERSION {
            return Err(VectorError::CorruptedIndex(format!(
                "unsupported version: expected {FORMAT_VERSION:#04x}, got {:#04x}",
                data[4]
            )));
        }

        // Read dimensions
        let dims = u32::from_le_bytes([data[5], data[6], data[7], data[8]]) as usize;
        if dims != config.dimensions {
            return Err(VectorError::CorruptedIndex(format!(
                "dimension mismatch: file has {dims}, config expects {}",
                config.dimensions
            )));
        }

        // Read count
        let count = u64::from_le_bytes([
            data[9], data[10], data[11], data[12], data[13], data[14], data[15], data[16],
        ]) as usize;

        // Validate total size
        let bytes_per_vector = 8 + dims * 4; // id (8) + floats (dims * 4)
        let expected_size = HEADER_SIZE + count * bytes_per_vector;
        if data.len() < expected_size {
            return Err(VectorError::CorruptedIndex(format!(
                "file truncated: expected at least {expected_size} bytes, got {}",
                data.len()
            )));
        }

        // Parse vectors
        let mut vectors = HashMap::with_capacity(count);
        let mut offset = HEADER_SIZE;
        for _ in 0..count {
            let id = u64::from_le_bytes([
                data[offset],
                data[offset + 1],
                data[offset + 2],
                data[offset + 3],
                data[offset + 4],
                data[offset + 5],
                data[offset + 6],
                data[offset + 7],
            ]);
            offset += 8;

            let mut vec = Vec::with_capacity(dims);
            for _ in 0..dims {
                let val = f32::from_le_bytes([
                    data[offset],
                    data[offset + 1],
                    data[offset + 2],
                    data[offset + 3],
                ]);
                vec.push(val);
                offset += 4;
            }

            vectors.insert(id, vec);
        }

        Ok(Self {
            vectors: RwLock::new(vectors),
            config: config.clone(),
        })
    }
}

#[cfg(test)]
mod tests;
