use crate::hash::Hash64;
use crate::model::ImageId;
use std::collections::hash_map::Entry;
use std::collections::HashMap;

/// A structure that finds image ids whose hash is within a Hamming radius.
pub trait CandidateIndex {
    fn query_within(&self, hash: Hash64, radius: u32) -> Vec<ImageId>;
}

struct Node {
    hash: Hash64,
    ids: Vec<ImageId>, // all images sharing exactly this hash
    children: HashMap<u32, Box<Node>>, // keyed by Hamming distance from this node
}

/// BK-tree over Hamming distance for radius queries without scanning the set.
pub struct BkTree {
    root: Option<Box<Node>>,
}

impl BkTree {
    pub fn new() -> Self {
        BkTree { root: None }
    }

    pub fn insert(&mut self, id: ImageId, hash: Hash64) {
        match &mut self.root {
            None => {
                self.root = Some(Box::new(Node { hash, ids: vec![id], children: HashMap::new() }));
            }
            Some(root) => {
                let mut node = root.as_mut();
                loop {
                    let d = node.hash.hamming(hash);
                    if d == 0 {
                        node.ids.push(id);
                        return;
                    }
                    match node.children.entry(d) {
                        Entry::Occupied(e) => node = e.into_mut(),
                        Entry::Vacant(e) => {
                            e.insert(Box::new(Node { hash, ids: vec![id], children: HashMap::new() }));
                            return;
                        }
                    }
                }
            }
        }
    }
}

impl Default for BkTree {
    fn default() -> Self {
        Self::new()
    }
}

impl CandidateIndex for BkTree {
    fn query_within(&self, hash: Hash64, radius: u32) -> Vec<ImageId> {
        let mut out = Vec::new();
        let mut stack: Vec<&Node> = Vec::new();
        if let Some(root) = &self.root {
            stack.push(root);
        }
        while let Some(node) = stack.pop() {
            let d = node.hash.hamming(hash);
            if d <= radius {
                out.extend_from_slice(&node.ids);
            }
            let low = d.saturating_sub(radius);
            let high = d + radius;
            for (k, child) in &node.children {
                if *k >= low && *k <= high {
                    stack.push(child);
                }
            }
        }
        out
    }
}

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

    fn brute_force(hashes: &[(usize, Hash64)], query: Hash64, radius: u32) -> Vec<usize> {
        let mut ids: Vec<usize> = hashes
            .iter()
            .filter(|(_, h)| h.hamming(query) <= radius)
            .map(|(id, _)| *id)
            .collect();
        ids.sort_unstable();
        ids
    }

    #[test]
    fn query_matches_brute_force() {
        let hashes = vec![
            (0, Hash64(0b0000)),
            (1, Hash64(0b0001)),
            (2, Hash64(0b0011)),
            (3, Hash64(0b1111)),
            (4, Hash64(0b0001)), // duplicate hash of id 1
        ];
        let mut tree = BkTree::new();
        for (id, h) in &hashes {
            tree.insert(*id, *h);
        }

        for &radius in &[0u32, 1, 2, 4] {
            for (_, q) in &hashes {
                let mut got = tree.query_within(*q, radius);
                got.sort_unstable();
                assert_eq!(got, brute_force(&hashes, *q, radius), "radius {radius}");
            }
        }
    }

    #[test]
    fn empty_tree_returns_nothing() {
        let tree = BkTree::new();
        assert!(tree.query_within(Hash64(0), 10).is_empty());
    }
}
