//! BK-tree implementation for visual similarity search.

use super::hamming::{hamming_distance, VisualIndex};
use crate::error::Result;
use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::Arc;
use stemedb_core::types::{Hash, PHash};
use tracing::{debug, instrument};

#[derive(Debug, Clone, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)]
#[archive(check_bytes)]
pub struct BkNode {
    /// The perceptual hash at this node.
    phash: PHash,
    /// The assertion hash (for result mapping).
    assertion_hash: Hash,
    /// Children indexed by hamming distance from this node's phash.
    /// Key = hamming distance, Value = child node index.
    children: HashMap<u32, usize>,
}

/// Snapshot of BK-tree state for persistence.
#[derive(Debug, Clone, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)]
#[archive(check_bytes)]
pub struct BkTreeSnapshot {
    /// All nodes in the tree.
    pub nodes: Vec<BkNode>,
    /// Maps assertion hashes to node indices.
    pub hash_to_node: HashMap<Hash, usize>,
    /// Total number of nodes in the tree.
    pub node_count: usize,
}

/// BK-tree based visual index implementation.
///
/// Provides O(log N) visual similarity search using hamming distance
/// as the metric space. The tree structure exploits the triangle
/// inequality to prune the search space.
pub struct BkTreeVisualIndex {
    /// All nodes in the tree (index 0 is root if non-empty).
    nodes: RwLock<Vec<BkNode>>,
    /// Maps assertion hashes to node indices (for deduplication).
    hash_to_node: RwLock<HashMap<Hash, usize>>,
}

impl BkTreeVisualIndex {
    /// Create a new empty visual index.
    pub fn new() -> Self {
        Self { nodes: RwLock::new(Vec::new()), hash_to_node: RwLock::new(HashMap::new()) }
    }

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

    /// Create a snapshot of the current index state for persistence.
    pub fn snapshot(&self) -> Result<BkTreeSnapshot> {
        let nodes = self.nodes.read().clone();
        let hash_to_node = self.hash_to_node.read().clone();
        let node_count = nodes.len();
        Ok(BkTreeSnapshot { nodes, hash_to_node, node_count })
    }

    /// Restore index from a snapshot.
    pub fn from_snapshot(snapshot: BkTreeSnapshot) -> Result<Self> {
        Ok(Self {
            nodes: RwLock::new(snapshot.nodes),
            hash_to_node: RwLock::new(snapshot.hash_to_node),
        })
    }

    /// Recursive search helper.
    fn search_recursive(
        nodes: &[BkNode],
        node_idx: usize,
        query: &PHash,
        threshold: u32,
        results: &mut Vec<(Hash, u32)>,
    ) {
        let node = &nodes[node_idx];
        let distance = hamming_distance(&node.phash, query);

        // If within threshold, add to results
        if distance <= threshold {
            results.push((node.assertion_hash, distance));
        }

        // Explore children with edge distance d where |d - distance| <= threshold
        // This is the key optimization: we only visit children that could
        // potentially have nodes within threshold of the query.
        let min_edge = distance.saturating_sub(threshold);
        let max_edge = distance.saturating_add(threshold);

        for (&edge_distance, &child_idx) in &node.children {
            if edge_distance >= min_edge && edge_distance <= max_edge {
                Self::search_recursive(nodes, child_idx, query, threshold, results);
            }
        }
    }
}

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

impl VisualIndex for BkTreeVisualIndex {
    #[instrument(skip(self, phash), fields(hash = %hex::encode(hash), phash = %hex::encode(phash)))]
    fn insert(&self, hash: &Hash, phash: &PHash) -> Result<()> {
        // Check if already indexed (idempotency)
        {
            let hash_to_node = self.hash_to_node.read();
            if hash_to_node.contains_key(hash) {
                debug!(hash = %hex::encode(hash), "Visual hash already indexed, skipping");
                return Ok(());
            }
        }

        let mut nodes = self.nodes.write();
        let mut hash_to_node = self.hash_to_node.write();

        let new_node = BkNode { phash: *phash, assertion_hash: *hash, children: HashMap::new() };

        if nodes.is_empty() {
            // First node becomes root
            let node_idx = 0;
            nodes.push(new_node);
            hash_to_node.insert(*hash, node_idx);
            debug!(hash = %hex::encode(hash), "Inserted as root node");
            return Ok(());
        }

        // Walk down tree to find insertion point
        let new_idx = nodes.len();
        let mut current_idx = 0;

        loop {
            let distance = hamming_distance(&nodes[current_idx].phash, phash);

            if let Some(&child_idx) = nodes[current_idx].children.get(&distance) {
                // Child exists at this distance, continue down
                current_idx = child_idx;
            } else {
                // No child at this distance, insert here
                nodes.push(new_node);
                nodes[current_idx].children.insert(distance, new_idx);
                hash_to_node.insert(*hash, new_idx);
                debug!(
                    hash = %hex::encode(hash),
                    parent_idx = current_idx,
                    distance,
                    "Inserted into BK-tree"
                );
                return Ok(());
            }
        }
    }

    #[instrument(skip(self, query), fields(phash = %hex::encode(query), threshold = threshold))]
    fn search(&self, query: &PHash, threshold: u32) -> Result<Vec<(Hash, u32)>> {
        // Clamp threshold to valid range
        let threshold = threshold.min(64);

        let nodes = self.nodes.read();

        if nodes.is_empty() {
            debug!("Visual index is empty, returning no results");
            return Ok(Vec::new());
        }

        let mut results = Vec::new();
        Self::search_recursive(&nodes, 0, query, threshold, &mut results);

        // Sort by distance ascending
        results.sort_by_key(|(_, d)| *d);

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

        Ok(results)
    }

    fn len(&self) -> Result<usize> {
        Ok(self.hash_to_node.read().len())
    }

    fn checkpoint(&self) -> Result<()> {
        // In-memory index, no checkpoint needed
        Ok(())
    }
}

// Implement for Arc<T> to enable sharing
impl<T: VisualIndex> VisualIndex for Arc<T> {
    fn insert(&self, hash: &Hash, phash: &PHash) -> Result<()> {
        (**self).insert(hash, phash)
    }

    fn search(&self, query: &PHash, threshold: u32) -> Result<Vec<(Hash, u32)>> {
        (**self).search(query, threshold)
    }

    fn len(&self) -> Result<usize> {
        (**self).len()
    }

    fn checkpoint(&self) -> Result<()> {
        (**self).checkpoint()
    }
}

// ────────────────────────────────────────────────────────────────────────────
// Persistent Visual Index
