//! Candidate retrieval from indexes.
//!
//! This module handles fetching assertions from different indexes:
//! - Subject index (S:{subject})
//! - Compound index (SP:{subject}:{predicate})
//! - Vector index (HNSW k-NN)
//! - Visual index (BK-tree hamming distance)
//! - Full scan (H: prefix)

use std::sync::Arc;

use stemedb_core::types::Assertion;
use stemedb_storage::{key_codec, IndexStore, KVStore, VectorIndex, VisualIndex};
use tracing::debug;

use crate::error::{QueryError, Result};
use crate::query::parse_hex_phash;

use super::QueryEngine;

impl<S: KVStore + 'static> QueryEngine<S> {
    /// Fetch assertions for a specific subject using the subject index.
    pub(super) async fn fetch_by_subject(&self, subject: &str) -> Result<Vec<Assertion>> {
        let hash_list = self.index_store.get_by_subject(subject).await?;

        let mut results = Vec::with_capacity(hash_list.len());
        for hash in hash_list {
            let assertion_key = key_codec::assertion_key(subject, &hex::encode(hash));
            if let Some(data) = self.store.get(&assertion_key).await? {
                match self.deserialize_assertion(&data) {
                    Ok(assertion) => results.push(assertion),
                    Err(e) => {
                        debug!(hash = %hex::encode(hash), "Skipping malformed assertion: {:?}", e);
                    }
                }
            }
        }

        Ok(results)
    }

    /// Fetch assertions for a specific subject and predicate using the compound index.
    pub(super) async fn fetch_by_subject_predicate(
        &self,
        subject: &str,
        predicate: &str,
    ) -> Result<Vec<Assertion>> {
        let hash_list = self.index_store.get_by_subject_predicate(subject, predicate).await?;

        let mut results = Vec::with_capacity(hash_list.len());
        for hash in hash_list {
            let assertion_key = key_codec::assertion_key(subject, &hex::encode(hash));
            if let Some(data) = self.store.get(&assertion_key).await? {
                match self.deserialize_assertion(&data) {
                    Ok(assertion) => results.push(assertion),
                    Err(e) => {
                        debug!(hash = %hex::encode(hash), "Skipping malformed assertion: {:?}", e);
                    }
                }
            }
        }

        Ok(results)
    }

    /// Fetch assertions for multiple subjects, deduplicating by hash.
    ///
    /// Used for alias resolution where a single query subject expands to
    /// multiple aliased subjects (e.g., code:// and rfc:// paths).
    pub(super) async fn fetch_by_subjects(&self, subjects: &[String]) -> Result<Vec<Assertion>> {
        use std::collections::HashSet;
        let mut seen_hashes: HashSet<[u8; 32]> = HashSet::new();
        let mut results = Vec::new();

        for subject in subjects {
            let hash_list = self.index_store.get_by_subject(subject).await?;
            for hash in hash_list {
                if !seen_hashes.insert(hash) {
                    continue; // Already seen this hash from another subject
                }

                let assertion_key = key_codec::assertion_key(subject, &hex::encode(hash));
                if let Some(data) = self.store.get(&assertion_key).await? {
                    match self.deserialize_assertion(&data) {
                        Ok(assertion) => results.push(assertion),
                        Err(e) => {
                            debug!(hash = %hex::encode(hash), "Skipping malformed assertion: {:?}", e);
                        }
                    }
                }
            }
        }

        debug!(
            subjects_count = subjects.len(),
            total_assertions = results.len(),
            "Fetched assertions for multiple subjects"
        );
        Ok(results)
    }

    /// Fetch assertions for multiple subjects with predicate filter, deduplicating by hash.
    ///
    /// Used for alias resolution when both subject and predicate are specified.
    pub(super) async fn fetch_by_subjects_predicate(
        &self,
        subjects: &[String],
        predicate: &str,
    ) -> Result<Vec<Assertion>> {
        use std::collections::HashSet;
        let mut seen_hashes: HashSet<[u8; 32]> = HashSet::new();
        let mut results = Vec::new();

        for subject in subjects {
            let hash_list = self.index_store.get_by_subject_predicate(subject, predicate).await?;
            for hash in hash_list {
                if !seen_hashes.insert(hash) {
                    continue; // Already seen this hash from another subject
                }

                let assertion_key = key_codec::assertion_key(subject, &hex::encode(hash));
                if let Some(data) = self.store.get(&assertion_key).await? {
                    match self.deserialize_assertion(&data) {
                        Ok(assertion) => results.push(assertion),
                        Err(e) => {
                            debug!(hash = %hex::encode(hash), "Skipping malformed assertion: {:?}", e);
                        }
                    }
                }
            }
        }

        debug!(
            subjects_count = subjects.len(),
            predicate,
            total_assertions = results.len(),
            "Fetched assertions for multiple subjects with predicate"
        );
        Ok(results)
    }

    /// Fetch all assertions by scanning the subjects discovery index.
    ///
    /// This scans `\x00SUBJECTS:` to discover all known subjects, then fetches
    /// all assertions per subject. This is O(n) and should be avoided for large databases.
    /// Use subject/predicate indexes when possible.
    pub(super) async fn fetch_all_assertions(&self) -> Result<Vec<Assertion>> {
        // Discover all subjects via the subjects index
        let subject_entries = self.store.scan_prefix(&key_codec::subjects_scan_prefix()).await?;

        let mut assertions = Vec::new();
        for (key, _) in subject_entries {
            // Extract subject from key: \x00SUBJECTS:{subject}
            let subject = match key_codec::extract_subject_from_subjects_key(&key) {
                Some(s) => s,
                None => continue,
            };

            // Fetch all assertions for this subject via the subject index
            let hash_list = self.index_store.get_by_subject(&subject).await?;
            for hash in hash_list {
                let assertion_key = key_codec::assertion_key(&subject, &hex::encode(hash));
                if let Some(data) = self.store.get(&assertion_key).await? {
                    match self.deserialize_assertion(&data) {
                        Ok(assertion) => assertions.push(assertion),
                        Err(e) => {
                            debug!("Skipping malformed assertion: {:?}", e);
                        }
                    }
                }
            }
        }

        Ok(assertions)
    }

    /// Fetch assertions by vector similarity using the HNSW index.
    ///
    /// Returns the k nearest neighbors to the query vector, ordered by distance.
    pub(super) async fn fetch_by_vector_similarity(
        &self,
        vector_index: &Arc<dyn VectorIndex>,
        query_vector: &[f32],
        k: usize,
    ) -> Result<Vec<Assertion>> {
        // Use the vector index to find nearest neighbors
        let neighbors =
            vector_index.search(query_vector, k).map_err(|e| QueryError::Index(e.to_string()))?;

        debug!(candidates_count = neighbors.len(), "Vector index returned candidates");

        // Fetch assertions by their hashes using reverse index for subject lookup
        let mut results = Vec::with_capacity(neighbors.len());
        for (hash, distance) in neighbors {
            let hash_hex = hex::encode(hash);
            // Look up subject from reverse index
            let reverse_key = key_codec::hash_subject_key(&hash_hex);
            let subject = match self.store.get(&reverse_key).await? {
                Some(bytes) => match String::from_utf8(bytes) {
                    Ok(s) => s,
                    Err(_) => {
                        debug!(hash = %hash_hex, "Invalid UTF-8 in reverse index, skipping");
                        continue;
                    }
                },
                None => {
                    debug!(hash = %hash_hex, "No reverse index entry, skipping");
                    continue;
                }
            };

            let assertion_key = key_codec::assertion_key(&subject, &hash_hex);
            if let Some(data) = self.store.get(&assertion_key).await? {
                match self.deserialize_assertion(&data) {
                    Ok(assertion) => {
                        debug!(
                            hash = %hash_hex,
                            distance,
                            "Found assertion via vector index"
                        );
                        results.push(assertion);
                    }
                    Err(e) => {
                        debug!(hash = %hash_hex, "Skipping malformed assertion: {:?}", e);
                    }
                }
            }
        }

        Ok(results)
    }

    /// Fetch assertions by visual similarity using the BK-tree index.
    ///
    /// Returns assertions whose visual hash is within the threshold hamming distance.
    pub(super) async fn fetch_by_visual_similarity(
        &self,
        visual_index: &Arc<dyn VisualIndex>,
        query_hash_hex: &str,
        threshold: u32,
    ) -> Result<Vec<Assertion>> {
        // Parse the hex string to a PHash
        let query_hash = parse_hex_phash(query_hash_hex).ok_or_else(|| {
            QueryError::InvalidInput(format!(
                "Invalid visual hash hex string: {} (expected 16 hex characters)",
                query_hash_hex
            ))
        })?;

        // Use the visual index to find matches within threshold
        let matches = visual_index
            .search(&query_hash, threshold)
            .map_err(|e| QueryError::Index(e.to_string()))?;

        debug!(candidates_count = matches.len(), threshold, "Visual index returned candidates");

        // Fetch assertions by their hashes using reverse index for subject lookup
        let mut results = Vec::with_capacity(matches.len());
        for (hash, distance) in matches {
            let hash_hex = hex::encode(hash);
            // Look up subject from reverse index
            let reverse_key = key_codec::hash_subject_key(&hash_hex);
            let subject = match self.store.get(&reverse_key).await? {
                Some(bytes) => match String::from_utf8(bytes) {
                    Ok(s) => s,
                    Err(_) => {
                        debug!(hash = %hash_hex, "Invalid UTF-8 in reverse index, skipping");
                        continue;
                    }
                },
                None => {
                    debug!(hash = %hash_hex, "No reverse index entry, skipping");
                    continue;
                }
            };

            let assertion_key = key_codec::assertion_key(&subject, &hash_hex);
            if let Some(data) = self.store.get(&assertion_key).await? {
                match self.deserialize_assertion(&data) {
                    Ok(assertion) => {
                        debug!(
                            hash = %hash_hex,
                            distance,
                            "Found assertion via visual index"
                        );
                        results.push(assertion);
                    }
                    Err(e) => {
                        debug!(hash = %hash_hex, "Skipping malformed assertion: {:?}", e);
                    }
                }
            }
        }

        Ok(results)
    }
}
