use crate::{Metric, SearchResult};
use super::filter::Allowlist;
use std::collections::BinaryHeap;
use std::cmp::Ordering;

/// A flat, contiguous store of bit-packed quantized vectors.
/// Memory layout: [vec_0 | vec_1 | ... | vec_n], each `packed_bytes` long.
pub struct VectorStore {
    pub data:         Vec<u8>,
    pub ids:          Vec<u64>,
    /// Per-vector dequantized L2 norm for length-renormalized scoring.
    /// Empty when loaded from a v1 file (norm correction skipped).
    pub norms:        Vec<f32>,
    pub packed_bytes: usize,
}

impl VectorStore {
    pub fn new(packed_bytes: usize) -> Self {
        Self { data: Vec::new(), ids: Vec::new(), norms: Vec::new(), packed_bytes }
    }

    pub fn add(&mut self, id: u64, packed: &[u8], norm: f32) {
        debug_assert_eq!(packed.len(), self.packed_bytes);
        self.ids.push(id);
        self.data.extend_from_slice(packed);
        self.norms.push(norm);
    }

    pub fn len(&self)      -> usize { self.ids.len() }
    pub fn is_empty(&self) -> bool  { self.ids.is_empty() }

    /// Iterator over (id, packed_bytes_slice) pairs.
    pub fn iter(&self) -> impl Iterator<Item = (u64, &[u8])> {
        self.ids.iter().copied().zip(
            self.data.chunks_exact(self.packed_bytes)
        )
    }

    /// Raw byte slice for vector at position `pos`.
    #[inline]
    pub fn get_packed(&self, pos: usize) -> &[u8] {
        let start = pos * self.packed_bytes;
        &self.data[start..start + self.packed_bytes]
    }
}

/// Scored vector entry for the min-heap (lower score = lower priority).
/// We negate to turn BinaryHeap (max-heap) into a min-heap for Top-K.
#[derive(PartialEq)]
struct Candidate {
    score: f32,
    id:    u64,
}

impl Eq for Candidate {}

impl PartialOrd for Candidate {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Candidate {
    fn cmp(&self, other: &Self) -> Ordering {
        // Invert for min-heap: lower score should be evicted first.
        other.score.partial_cmp(&self.score)
            .unwrap_or(Ordering::Equal)
            .then(self.id.cmp(&other.id))
    }
}

/// Brute-force scan over the vector store.
///
/// Applies `score_fn` to every (allowed) candidate and returns the top-K.
/// `score_fn(packed_vector) -> f32` — higher is better for all metrics.
pub fn top_k<F>(
    store:      &VectorStore,
    k:          usize,
    allowlist:  Option<&Allowlist>,
    metric:     Metric,
    score_fn:   F,
) -> Vec<SearchResult>
where
    F: Fn(&[u8]) -> f32,
{
    if store.is_empty() || k == 0 {
        return Vec::new();
    }

    // Min-heap of capacity k: evicts the lowest score when full.
    let mut heap: BinaryHeap<Candidate> = BinaryHeap::with_capacity(k + 1);

    for (pos, (id, packed)) in store.iter().enumerate() {
        if let Some(al) = allowlist {
            if !al.contains(id) {
                continue;
            }
        }

        let raw   = score_fn(packed);
        let score = if store.norms.is_empty() {
            raw
        } else {
            metric.adjust_score(raw, store.norms[pos])
        };

        if heap.len() < k {
            heap.push(Candidate { score, id });
        } else if let Some(min) = heap.peek() {
            if score > min.score {
                heap.pop();
                heap.push(Candidate { score, id });
            }
        }
    }

    // Drain and sort descending (best first).
    let mut results: Vec<SearchResult> = heap
        .into_iter()
        .map(|c| SearchResult { id: c.id, score: c.score })
        .collect();
    results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(Ordering::Equal));
    results
}

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

    fn dummy_store(n: usize, packed_bytes: usize) -> VectorStore {
        let mut store = VectorStore::new(packed_bytes);
        for i in 0..n {
            let packed = vec![i as u8; packed_bytes];
            store.add(i as u64, &packed, 1.0);
        }
        store
    }

    #[test]
    fn returns_k_results() {
        let store = dummy_store(20, 4);
        let results = top_k(&store, 5, None, Metric::Cosine, |_| 1.0);
        assert_eq!(results.len(), 5);
    }

    #[test]
    fn fewer_than_k_vectors() {
        let store = dummy_store(3, 4);
        let results = top_k(&store, 10, None, Metric::Cosine, |_| 1.0);
        assert_eq!(results.len(), 3);
    }

    #[test]
    fn results_ordered_descending() {
        let store = dummy_store(10, 4);
        // Score = first byte of packed vector (0..9)
        let results = top_k(&store, 5, None, Metric::Cosine, |packed| packed[0] as f32);
        // Expect top 5: ids 9,8,7,6,5 (scores 9,8,7,6,5)
        assert_eq!(results[0].id, 9);
        assert_eq!(results[1].id, 8);
        for i in 1..results.len() {
            assert!(results[i-1].score >= results[i].score);
        }
    }

    #[test]
    fn allowlist_filters_correctly() {
        let store = dummy_store(10, 4);
        let al    = Allowlist::new(&[2, 5, 8]);
        let results = top_k(&store, 10, Some(&al), Metric::Cosine, |_| 1.0);
        assert_eq!(results.len(), 3);
        let ids: Vec<u64> = results.iter().map(|r| r.id).collect();
        assert!(ids.contains(&2));
        assert!(ids.contains(&5));
        assert!(ids.contains(&8));
    }

    #[test]
    fn empty_store() {
        let store   = VectorStore::new(4);
        let results = top_k(&store, 5, None, Metric::Cosine, |_| 0.0);
        assert!(results.is_empty());
    }
}
