terraphim_automata 1.11.0

use crate::{Result, TerraphimAutomataError};
use ahash::AHashMap;
use fst::{Automaton, IntoStreamer, Map, MapBuilder, Streamer, automaton::Str};
use serde::{Deserialize, Serialize};
use terraphim_types::{NormalizedTermValue, Thesaurus};

#[cfg(feature = "remote-loading")]
use crate::{AutomataPath, load_thesaurus};

/// Autocomplete index built from thesaurus data using FST
#[derive(Debug, Clone)]
pub struct AutocompleteIndex {
    /// FST for fast prefix searches
    fst: Map<Vec<u8>>,
    /// Metadata lookup for results
    metadata: AHashMap<String, AutocompleteMetadata>,
    /// Original thesaurus name
    name: String,
}

/// Metadata associated with each autocomplete term
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AutocompleteMetadata {
    pub id: u64,
    pub normalized_term: NormalizedTermValue,
    pub url: Option<String>,
    pub original_term: String,
}

/// Result from autocomplete search
#[derive(Debug, Clone, PartialEq)]
pub struct AutocompleteResult {
    pub term: String,
    pub normalized_term: NormalizedTermValue,
    pub id: u64,
    pub url: Option<String>,
    pub score: f64, // FST value as relevance score
}

/// Configuration for autocomplete behavior
#[derive(Debug, Clone)]
pub struct AutocompleteConfig {
    pub max_results: usize,
    pub min_prefix_length: usize,
    pub case_sensitive: bool,
}

impl Default for AutocompleteConfig {
    fn default() -> Self {
        Self {
            max_results: 10,
            min_prefix_length: 1,
            case_sensitive: false,
        }
    }
}

impl AutocompleteIndex {
    /// Get the name of the autocomplete index
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Get the number of terms in the autocomplete index
    pub fn len(&self) -> usize {
        self.metadata.len()
    }

    /// Check if the autocomplete index is empty
    pub fn is_empty(&self) -> bool {
        self.metadata.is_empty()
    }

    /// Iterate over metadata entries (term -> metadata)
    pub fn metadata_iter(&self) -> impl Iterator<Item = (&str, &AutocompleteMetadata)> {
        self.metadata.iter().map(|(k, v)| (k.as_str(), v))
    }

    /// Lookup metadata by term key (lowercased key as stored in index)
    pub fn metadata_get(&self, term: &str) -> Option<&AutocompleteMetadata> {
        self.metadata.get(term)
    }
}

/// Build autocomplete index from existing thesaurus
pub fn build_autocomplete_index(
    thesaurus: Thesaurus,
    config: Option<AutocompleteConfig>,
) -> Result<AutocompleteIndex> {
    let config = config.unwrap_or_default();
    let mut terms_with_scores: Vec<(String, u64)> = Vec::new();
    let mut metadata: AHashMap<String, AutocompleteMetadata> = AHashMap::new();

    log::debug!(
        "Building autocomplete index from thesaurus with {} entries",
        thesaurus.len()
    );

    // Extract all terms from thesaurus and assign scores based on term frequency/importance
    for (key, normalized_term) in &thesaurus {
        let term = if config.case_sensitive {
            key.to_string()
        } else {
            key.as_str().to_lowercase()
        };

        // Use the ID as a score (higher ID = higher relevance)
        // In a real implementation, this could be based on term frequency, importance, etc.
        let score = normalized_term.id;

        terms_with_scores.push((term.clone(), score));

        metadata.insert(
            term.clone(),
            AutocompleteMetadata {
                id: normalized_term.id,
                normalized_term: normalized_term.value.clone(),
                url: normalized_term.url.clone(),
                original_term: if config.case_sensitive {
                    key.to_string()
                } else {
                    key.as_str().to_lowercase()
                },
            },
        );
    }

    // Sort terms lexicographically for FST building
    terms_with_scores.sort_by(|a, b| a.0.cmp(&b.0));

    log::debug!("Building FST with {} sorted terms", terms_with_scores.len());

    // Build FST
    let mut builder = MapBuilder::memory();
    for (term, score) in terms_with_scores {
        builder.insert(&term, score)?;
    }

    let fst_bytes = builder.into_inner()?;
    let fst = Map::new(fst_bytes)?;

    log::debug!(
        "Successfully built autocomplete index with {} terms",
        metadata.len()
    );

    Ok(AutocompleteIndex {
        fst,
        metadata,
        name: thesaurus.name().to_string(),
    })
}

/// Load thesaurus and build autocomplete index in one step
///
/// Note: This function requires the "remote-loading" feature to be enabled
/// for async loading of remote thesaurus files.
#[cfg(feature = "remote-loading")]
pub async fn load_autocomplete_index(
    automata_path: &AutomataPath,
    config: Option<AutocompleteConfig>,
) -> Result<AutocompleteIndex> {
    log::debug!("Loading thesaurus from: {}", automata_path);
    let thesaurus = load_thesaurus(automata_path).await?;
    build_autocomplete_index(thesaurus, config)
}

/// Perform autocomplete search with prefix
pub fn autocomplete_search(
    index: &AutocompleteIndex,
    prefix: &str,
    limit: Option<usize>,
) -> Result<Vec<AutocompleteResult>> {
    let config = AutocompleteConfig::default();
    let search_prefix = if config.case_sensitive {
        prefix.to_string()
    } else {
        prefix.to_lowercase()
    };

    if search_prefix.len() < config.min_prefix_length {
        return Ok(Vec::new());
    }

    let max_results = limit.unwrap_or(config.max_results);
    let mut results = Vec::new();

    log::trace!(
        "Searching autocomplete index for prefix: '{}'",
        search_prefix
    );

    // Use FST to find all terms with the given prefix using the Str automaton
    let automaton = Str::new(&search_prefix).starts_with();
    let mut stream = index.fst.search(automaton).into_stream();

    while let Some((term_bytes, score)) = stream.next() {
        if results.len() >= max_results {
            break;
        }

        let term = String::from_utf8_lossy(term_bytes).to_string();

        if let Some(metadata) = index.metadata.get(&term) {
            results.push(AutocompleteResult {
                term: metadata.original_term.clone(),
                normalized_term: metadata.normalized_term.clone(),
                id: metadata.id,
                url: metadata.url.clone(),
                score: score as f64,
            });
        }
    }

    // Sort results by score (descending) then by term length (ascending) for better UX
    results.sort_by(|a, b| {
        b.score
            .partial_cmp(&a.score)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then_with(|| a.term.len().cmp(&b.term.len()))
    });

    log::trace!(
        "Found {} autocomplete results for prefix: '{}'",
        results.len(),
        search_prefix
    );

    Ok(results)
}

/// Fuzzy autocomplete search using Levenshtein edit distance (baseline comparison)
///
/// Uses Levenshtein distance calculation for baseline comparison with the default
/// Jaro-Winkler fuzzy search. Levenshtein is useful when you need exact edit distance control.
pub fn fuzzy_autocomplete_search_levenshtein(
    index: &AutocompleteIndex,
    prefix: &str,
    max_edit_distance: usize,
    limit: Option<usize>,
) -> Result<Vec<AutocompleteResult>> {
    let max_results = limit.unwrap_or(10);
    let mut all_results = Vec::new();

    // Try exact prefix first
    let exact_results = autocomplete_search(index, prefix, Some(max_results))?;
    all_results.extend(exact_results);

    if all_results.len() >= max_results {
        all_results.truncate(max_results);
        return Ok(all_results);
    }

    // For fuzzy matching, scan all terms and calculate Levenshtein similarity
    if max_edit_distance > 0 {
        let mut fuzzy_candidates = Vec::new();

        // Iterate through all terms in the metadata to find fuzzy matches
        for (term, metadata) in &index.metadata {
            // Skip if we already have this result from exact search
            if all_results.iter().any(|r| r.id == metadata.id) {
                continue;
            }

            // Calculate Levenshtein distance - check both full term and individual words
            let distances = {
                let mut dists = vec![strsim::levenshtein(prefix, term)];

                // Also check against individual words in the term for better fuzzy matching
                for word in term.split_whitespace() {
                    dists.push(strsim::levenshtein(prefix, word));
                }

                dists
            };

            let min_distance = distances.into_iter().min().unwrap_or(usize::MAX);

            // Only include if within edit distance threshold
            if min_distance <= max_edit_distance {
                // Convert distance to similarity score (same as terraphim_service scorer)
                let similarity = 1.0 / (1.0 + min_distance as f64);

                // Weight by original FST score
                let original_score = metadata.id as f64;
                let combined_score = similarity * original_score * 0.8; // Penalize fuzzy matches

                fuzzy_candidates.push(AutocompleteResult {
                    term: metadata.original_term.clone(),
                    normalized_term: metadata.normalized_term.clone(),
                    id: metadata.id,
                    url: metadata.url.clone(),
                    score: combined_score,
                });
            }
        }

        // Sort by combined score (similarity * original_score)
        fuzzy_candidates.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| a.term.len().cmp(&b.term.len()))
        });

        // Add the best fuzzy matches
        let remaining_slots = max_results - all_results.len();
        all_results.extend(fuzzy_candidates.into_iter().take(remaining_slots));
    }

    // Final sort by score
    all_results.sort_by(|a, b| {
        b.score
            .partial_cmp(&a.score)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then_with(|| a.term.len().cmp(&b.term.len()))
    });

    all_results.truncate(max_results);
    Ok(all_results)
}

/// Fuzzy autocomplete search using Jaro-Winkler similarity (DEFAULT)
///
/// Jaro-Winkler is the recommended algorithm for autocomplete because it gives extra weight
/// to common prefixes and handles character transpositions better. It's 2.3x faster than
/// Levenshtein and produces higher quality results for autocomplete scenarios.
pub fn fuzzy_autocomplete_search(
    index: &AutocompleteIndex,
    prefix: &str,
    min_similarity: f64,
    limit: Option<usize>,
) -> Result<Vec<AutocompleteResult>> {
    let max_results = limit.unwrap_or(10);
    let mut all_results = Vec::new();

    // Try exact prefix first
    let exact_results = autocomplete_search(index, prefix, Some(max_results))?;
    all_results.extend(exact_results);

    if all_results.len() >= max_results {
        all_results.truncate(max_results);
        return Ok(all_results);
    }

    // For fuzzy matching, scan all terms and calculate Jaro-Winkler similarity
    if min_similarity > 0.0 && min_similarity < 1.0 {
        let mut fuzzy_candidates = Vec::new();

        // Iterate through all terms in the metadata to find fuzzy matches
        for (term, metadata) in &index.metadata {
            // Skip if we already have this result from exact search
            if all_results.iter().any(|r| r.id == metadata.id) {
                continue;
            }

            // Calculate Jaro-Winkler similarity - check both full term and individual words
            let similarities = {
                let mut sims = vec![strsim::jaro_winkler(prefix, term)];

                // Also check against individual words in the term for better fuzzy matching
                for word in term.split_whitespace() {
                    sims.push(strsim::jaro_winkler(prefix, word));
                }

                sims
            };

            let max_similarity = similarities
                .into_iter()
                .fold(0.0f64, |acc, sim| acc.max(sim));

            // Only include if above similarity threshold
            if max_similarity >= min_similarity {
                // Weight by original FST score
                let original_score = metadata.id as f64;
                let combined_score = max_similarity * original_score * 0.8; // Penalize fuzzy matches

                fuzzy_candidates.push(AutocompleteResult {
                    term: metadata.original_term.clone(),
                    normalized_term: metadata.normalized_term.clone(),
                    id: metadata.id,
                    url: metadata.url.clone(),
                    score: combined_score,
                });
            }
        }

        // Sort by combined score (similarity * original_score)
        fuzzy_candidates.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| a.term.len().cmp(&b.term.len()))
        });

        // Add the best fuzzy matches
        let remaining_slots = max_results - all_results.len();
        all_results.extend(fuzzy_candidates.into_iter().take(remaining_slots));
    }

    // Final sort by score
    all_results.sort_by(|a, b| {
        b.score
            .partial_cmp(&a.score)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then_with(|| a.term.len().cmp(&b.term.len()))
    });

    all_results.truncate(max_results);
    Ok(all_results)
}

/// Enhanced fuzzy autocomplete search using Jaro-Winkler similarity (DEPRECATED - use fuzzy_autocomplete_search)
///
/// This function is deprecated. Use `fuzzy_autocomplete_search` instead, which now uses
/// Jaro-Winkler as the default algorithm for better autocomplete performance.
#[deprecated(since = "0.1.0", note = "Use fuzzy_autocomplete_search instead")]
pub fn fuzzy_autocomplete_search_jaro_winkler(
    index: &AutocompleteIndex,
    prefix: &str,
    min_similarity: f64,
    limit: Option<usize>,
) -> Result<Vec<AutocompleteResult>> {
    fuzzy_autocomplete_search(index, prefix, min_similarity, limit)
}

/// Serialize index to bytes for caching
pub fn serialize_autocomplete_index(index: &AutocompleteIndex) -> Result<Vec<u8>> {
    // Create a serializable version of the index
    let serializable = SerializableIndex {
        fst_bytes: index.fst.as_fst().as_bytes().to_vec(),
        metadata: index.metadata.clone(),
        name: index.name.clone(),
    };

    bincode::serialize(&serializable)
        .map_err(|e| TerraphimAutomataError::Dict(format!("Serialization error: {}", e)))
}

/// Deserialize index from bytes
pub fn deserialize_autocomplete_index(data: &[u8]) -> Result<AutocompleteIndex> {
    let serializable: SerializableIndex = bincode::deserialize(data)
        .map_err(|e| TerraphimAutomataError::Dict(format!("Deserialization error: {}", e)))?;

    let fst = Map::new(serializable.fst_bytes)?;

    Ok(AutocompleteIndex {
        fst,
        metadata: serializable.metadata,
        name: serializable.name,
    })
}

/// Serializable version of AutocompleteIndex for persistence
#[derive(Serialize, Deserialize)]
struct SerializableIndex {
    fst_bytes: Vec<u8>,
    metadata: AHashMap<String, AutocompleteMetadata>,
    name: String,
}

#[cfg(test)]
mod tests {
    use super::*;
    use terraphim_types::{NormalizedTerm, NormalizedTermValue, Thesaurus};

    fn create_test_thesaurus() -> Thesaurus {
        let mut thesaurus = Thesaurus::new("Test".to_string());

        // Add test terms with different IDs (scores)
        let terms = vec![
            ("machine learning", "machine learning", 10),
            ("ml", "machine learning", 10),
            ("artificial intelligence", "artificial intelligence", 20),
            ("ai", "artificial intelligence", 20),
            ("neural network", "neural network", 15),
            ("deep learning", "deep learning", 18),
            ("data science", "data science", 12),
            ("algorithm", "algorithm", 8),
            ("programming", "programming", 5),
            ("python", "python", 7),
        ];

        for (key, normalized, id) in terms {
            let normalized_term = NormalizedTerm {
                id,
                value: NormalizedTermValue::from(normalized),
                display_value: None,
                url: Some(format!(
                    "https://example.com/{}",
                    normalized.replace(' ', "-")
                )),
            };
            thesaurus.insert(NormalizedTermValue::from(key), normalized_term);
        }

        thesaurus
    }

    #[test]
    fn test_build_autocomplete_index() {
        let thesaurus = create_test_thesaurus();
        let index = build_autocomplete_index(thesaurus, None).unwrap();

        assert_eq!(index.name(), "Test");
        assert_eq!(index.len(), 10);
        assert!(!index.is_empty());
    }

    #[test]
    fn test_autocomplete_search_basic() {
        let thesaurus = create_test_thesaurus();
        let index = build_autocomplete_index(thesaurus, None).unwrap();

        // Test prefix search
        let results = autocomplete_search(&index, "ma", None).unwrap();
        assert!(!results.is_empty());

        // Check that results contain "machine learning"
        let has_ml = results.iter().any(|r| r.term == "machine learning");
        assert!(has_ml, "Should find 'machine learning' for prefix 'ma'");

        // Test exact match
        let results = autocomplete_search(&index, "python", None).unwrap();
        assert!(!results.is_empty());
        assert!(results.iter().any(|r| r.term == "python"));
    }

    #[test]
    fn test_autocomplete_search_ordering() {
        let thesaurus = create_test_thesaurus();
        let index = build_autocomplete_index(thesaurus, None).unwrap();

        // Search for terms that should be ordered by score
        let results = autocomplete_search(&index, "a", Some(5)).unwrap();

        // Results should be sorted by score (descending)
        for i in 1..results.len() {
            assert!(
                results[i - 1].score >= results[i].score,
                "Results should be sorted by score (descending)"
            );
        }
    }

    #[test]
    fn test_autocomplete_search_limits() {
        let thesaurus = create_test_thesaurus();
        let index = build_autocomplete_index(thesaurus, None).unwrap();

        // Test limit parameter
        let results = autocomplete_search(&index, "", Some(3)).unwrap();
        assert!(results.len() <= 3, "Should respect limit parameter");

        // Test default limit
        let results = autocomplete_search(&index, "", None).unwrap();
        assert!(results.len() <= 10, "Should respect default limit");
    }

    #[test]
    fn test_fuzzy_autocomplete_search() {
        let thesaurus = create_test_thesaurus();
        let index = build_autocomplete_index(thesaurus, None).unwrap();

        // Test fuzzy search with Jaro-Winkler similarity (now the default)
        let results = fuzzy_autocomplete_search(&index, "machne", 0.6, Some(5)).unwrap();

        // Should find results even with typo - "machine" should match "machne" with good similarity
        assert!(
            !results.is_empty(),
            "Fuzzy search should find results for 'machne'"
        );

        // Verify we get machine learning with good similarity score
        let has_ml = results.iter().any(|r| r.term == "machine learning");
        assert!(
            has_ml,
            "Should find 'machine learning' for fuzzy search 'machne'"
        );
    }

    #[test]
    fn test_fuzzy_search_levenshtein_scoring() {
        let thesaurus = create_test_thesaurus();
        let index = build_autocomplete_index(thesaurus, None).unwrap();

        // Test different edit distances using Levenshtein baseline function
        let results_distance_1 =
            fuzzy_autocomplete_search_levenshtein(&index, "pythno", 1, Some(10)).unwrap();
        let results_distance_2 =
            fuzzy_autocomplete_search_levenshtein(&index, "pythno", 2, Some(10)).unwrap();

        // Should find more results with higher edit distance
        assert!(
            results_distance_2.len() >= results_distance_1.len(),
            "Higher edit distance should yield more or equal results"
        );

        // Test exact match should score higher than fuzzy match
        let exact_results = autocomplete_search(&index, "python", None).unwrap();
        let fuzzy_results = fuzzy_autocomplete_search(&index, "pythno", 0.6, None).unwrap();

        if !exact_results.is_empty() && !fuzzy_results.is_empty() {
            let exact_python = exact_results.iter().find(|r| r.term == "python");
            let fuzzy_python = fuzzy_results.iter().find(|r| r.term == "python");

            if let (Some(exact), Some(fuzzy)) = (exact_python, fuzzy_python) {
                // Exact match should have higher score than fuzzy match
                // (though fuzzy search applies penalty factor)
                assert!(
                    exact.score > fuzzy.score,
                    "Exact match should score higher than fuzzy match"
                );
            }
        }
    }

    #[test]
    fn test_serialization_roundtrip() {
        let thesaurus = create_test_thesaurus();
        let original_index = build_autocomplete_index(thesaurus, None).unwrap();

        // Serialize
        let serialized = serialize_autocomplete_index(&original_index).unwrap();
        assert!(
            !serialized.is_empty(),
            "Serialized data should not be empty"
        );

        // Deserialize
        let deserialized_index = deserialize_autocomplete_index(&serialized).unwrap();

        // Verify integrity
        assert_eq!(original_index.name(), deserialized_index.name());
        assert_eq!(original_index.len(), deserialized_index.len());

        // Test that search works the same
        let original_results = autocomplete_search(&original_index, "ma", None).unwrap();
        let deserialized_results = autocomplete_search(&deserialized_index, "ma", None).unwrap();

        assert_eq!(original_results.len(), deserialized_results.len());
        for (orig, deser) in original_results.iter().zip(deserialized_results.iter()) {
            assert_eq!(orig.term, deser.term);
            assert_eq!(orig.id, deser.id);
            assert_eq!(orig.score, deser.score);
        }
    }

    #[cfg(feature = "remote-loading")]
    #[tokio::test]
    async fn test_load_autocomplete_index() {
        // Test loading from local example
        let result = load_autocomplete_index(&AutomataPath::local_example(), None).await;

        match result {
            Ok(index) => {
                assert!(!index.is_empty(), "Loaded index should not be empty");
                assert_eq!(index.name(), "Engineering"); // From local example file

                // Test search functionality
                let results = autocomplete_search(&index, "foo", None).unwrap();
                assert!(
                    !results.is_empty(),
                    "Should find results for 'foo' in test data"
                );
            }
            Err(e) => {
                // This is acceptable if the test data file doesn't exist
                log::warn!("Could not load test data for autocomplete index: {}", e);
            }
        }
    }

    #[test]
    fn test_autocomplete_config() {
        let config = AutocompleteConfig {
            max_results: 3,
            min_prefix_length: 2,
            case_sensitive: false,
        };

        let thesaurus = create_test_thesaurus();
        let index = build_autocomplete_index(thesaurus, Some(config)).unwrap();

        // Test that short prefixes return no results
        let _results = autocomplete_search(&index, "a", None).unwrap();
        // Note: The config is only used during index building, not search
        // In a full implementation, we'd pass config to search function too
    }
}