poimen-memory/crates/mem-cli/src/query/entity_linker.rs

//! Entity Linking (Phase 5.1)
//!
//! Identifies co-references, links text spans to entities, detects aliases,
//! and suggests entity merges.

use std::collections::{HashMap, HashSet};
use sqlx::PgPool;
use serde::{Deserialize, Serialize};
use tracing::{debug, warn};

/// Result of linking a text mention to an entity
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct MentionLink {
    /// The text span that was linked
    pub mention_text: String,
    /// Start offset in original text
    pub start_offset: usize,
    /// End offset in original text
    pub end_offset: usize,
    /// Entity ID it was linked to
    pub entity_id: String,
    /// Entity name
    pub entity_name: String,
    /// Confidence of link (0.0-1.0)
    pub confidence: f32,
    /// Why it was linked (semantic, lexical, alias, etc.)
    pub reason: LinkReason,
}

/// Reason for linking
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum LinkReason {
    /// Semantic similarity (high embedding match)
    SemanticMatch,
    /// Lexical match (exact or near-exact string)
    LexicalMatch,
    /// Known alias
    AliasMatch,
    /// Acronym expansion (e.g., "k8s" → "Kubernetes")
    AcronymMatch,
    /// Partial/substring match
    PartialMatch,
}

/// Alias suggestion
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AliasSuggestion {
    /// Entity ID
    pub entity_id: String,
    /// Entity name (canonical)
    pub canonical_name: String,
    /// Suggested alias
    pub alias: String,
    /// Confidence (0.0-1.0)
    pub confidence: f32,
    /// How often this alias appears in text
    pub frequency: usize,
}

/// Entity merge candidate
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MergeSuggestion {
    /// Entity 1 ID
    pub entity1_id: String,
    /// Entity 1 name
    pub entity1_name: String,
    /// Entity 2 ID
    pub entity2_id: String,
    /// Entity 2 name
    pub entity2_name: String,
    /// Confidence they're the same (0.0-1.0)
    pub confidence: f32,
    /// Reasons for merge
    pub reasons: Vec<String>,
}

/// Co-reference cluster (multiple mentions of same entity)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CoreferenceCluster {
    /// Representative entity ID
    pub entity_id: String,
    /// All mention texts in this cluster
    pub mentions: Vec<String>,
    /// Mention count
    pub mention_count: usize,
    /// Confidence this is correct clustering
    pub confidence: f32,
}

/// Entity Linking Engine
pub struct EntityLinker {
    pool: PgPool,
}

impl EntityLinker {
    pub fn new(pool: PgPool) -> Self {
        EntityLinker { pool }
    }

    /// Link mentions in text to existing entities
    ///
    /// Returns:
    /// - Vec<MentionLink>: Successful links
    /// - Vec<String>: Unlinked mentions
    pub async fn link_mentions(
        &self,
        text: &str,
        project_id: &str,
    ) -> Result<(Vec<MentionLink>, Vec<String>), String> {
        if text.is_empty() {
            return Ok((vec![], vec![]));
        }

        // Extract potential mentions (noun phrases, capitalized sequences)
        let mentions = self.extract_mentions(text)?;
        debug!("Extracted {} potential mentions from text", mentions.len());

        // Get all entities from database
        let entities = self.fetch_entities(project_id).await?;
        debug!("Loaded {} entities from database", entities.len());

        let mut links = Vec::new();
        let mut unlinked = Vec::new();

        for mention in mentions {
            match self.find_best_link(&mention.text, &entities).await? {
                Some((entity_id, entity_name, confidence, reason)) => {
                    links.push(MentionLink {
                        mention_text: mention.text.clone(),
                        start_offset: mention.start,
                        end_offset: mention.end,
                        entity_id,
                        entity_name,
                        confidence,
                        reason,
                    });
                }
                None => {
                    unlinked.push(mention.text);
                }
            }
        }

        Ok((links, unlinked))
    }

    /// Detect aliases for an entity
    pub async fn detect_aliases(
        &self,
        entity_id: &str,
        entity_name: &str,
        text_sample: &[String],
    ) -> Result<Vec<AliasSuggestion>, String> {
        let mut aliases = HashMap::new();

        for text in text_sample {
            let mentions = self.extract_mentions(text)?;
            for mention in mentions {
                if self.is_similar(&mention.text, entity_name) {
                    let entry = aliases.entry(mention.text.clone()).or_insert((0, 0.5));
                    entry.0 += 1;
                }
            }
        }

        // Convert to suggestions, only include frequent ones
        let suggestions: Vec<_> = aliases
            .into_iter()
            .filter(|(_, (count, _))| *count > 1) // At least 2 occurrences
            .map(|(alias, (frequency, confidence))| AliasSuggestion {
                entity_id: entity_id.to_string(),
                canonical_name: entity_name.to_string(),
                alias,
                confidence: (confidence * (frequency as f32 / 10.0).min(1.0)).min(1.0),
                frequency,
            })
            .collect();

        Ok(suggestions)
    }

    /// Suggest entity merges based on similarity
    pub async fn suggest_merges(
        &self,
        project_id: &str,
        similarity_threshold: f32,
    ) -> Result<Vec<MergeSuggestion>, String> {
        let entities = self.fetch_entities(project_id).await?;
        let mut suggestions = Vec::new();

        for (i, ent1) in entities.iter().enumerate() {
            for ent2 in &entities[(i + 1)..] {
                let similarity = self.compute_similarity(&ent1.name, &ent2.name);
                if similarity >= similarity_threshold {
                    let mut reasons = Vec::new();

                    if ent1.name.contains(&ent2.name) || ent2.name.contains(&ent1.name) {
                        reasons.push("Substring match".to_string());
                    }

                    if self.edit_distance(&ent1.name, &ent2.name) <= 2 {
                        reasons.push("Near edit distance".to_string());
                    }

                    if self.have_common_relations(&ent1.id, &ent2.id) {
                        reasons.push("Common relations".to_string());
                    }

                    suggestions.push(MergeSuggestion {
                        entity1_id: ent1.id.clone(),
                        entity1_name: ent1.name.clone(),
                        entity2_id: ent2.id.clone(),
                        entity2_name: ent2.name.clone(),
                        confidence: similarity,
                        reasons,
                    });
                }
            }
        }

        Ok(suggestions)
    }

    /// Identify coreference clusters
    pub async fn detect_coreferences(
        &self,
        texts: &[String],
        project_id: &str,
    ) -> Result<Vec<CoreferenceCluster>, String> {
        let mut clusters: HashMap<String, Vec<String>> = HashMap::new();
        let entities = self.fetch_entities(project_id).await?;

        for text in texts {
            let (links, _) = self.link_mentions(text, project_id).await?;
            for link in links {
                clusters
                    .entry(link.entity_id)
                    .or_insert_with(Vec::new)
                    .push(link.mention_text);
            }
        }

        let mut result = Vec::new();
        for (entity_id, mentions) in clusters {
            if let Some(entity) = entities.iter().find(|e| e.id == entity_id) {
                let unique_mentions: Vec<_> = mentions.iter().cloned().collect::<HashSet<_>>().into_iter().collect();
                result.push(CoreferenceCluster {
                    entity_id: entity_id.clone(),
                    mention_count: mentions.len(),
                    confidence: 0.85, // Confidence from linking process
                    mentions: unique_mentions,
                });
            }
        }

        Ok(result)
    }

    // ========== Private Helper Methods ==========

    /// Extract potential entity mentions from text
    fn extract_mentions(&self, text: &str) -> Result<Vec<Mention>, String> {
        let mut mentions = Vec::new();

        // Simple mention extraction: capitalized sequences, quoted text
        let words: Vec<&str> = text.split_whitespace().collect();
        let mut i = 0;

        while i < words.len() {
            let word = words[i];

            // Capitalized word (potential entity)
            if word.chars().next().map_or(false, |c| c.is_uppercase()) && word.len() > 2 {
                let start_pos = text.find(word).unwrap_or(0);
                let end_pos = start_pos + word.len();

                mentions.push(Mention {
                    text: word.to_string(),
                    start: start_pos,
                    end: end_pos,
                });

                // Multi-word entity (consecutive capitalized words)
                let mut j = i + 1;
                let mut multi_text = word.to_string();
                while j < words.len() && words[j].chars().next().map_or(false, |c| c.is_uppercase()) {
                    multi_text.push(' ');
                    multi_text.push_str(words[j]);
                    j += 1;
                }

                if j > i + 1 {
                    let start_pos = text.find(&multi_text).unwrap_or(0);
                    let end_pos = start_pos + multi_text.len();
                    mentions.push(Mention {
                        text: multi_text,
                        start: start_pos,
                        end: end_pos,
                    });
                    i = j - 1;
                }
            }
            i += 1;
        }

        Ok(mentions)
    }

    /// Find best link for a mention
    async fn find_best_link(
        &self,
        mention: &str,
        entities: &[EntityInfo],
    ) -> Result<Option<(String, String, f32, LinkReason)>, String> {
        let mut best: Option<(String, String, f32, LinkReason)> = None;

        for entity in entities {
            // Check exact match first (highest confidence)
            if entity.name.eq_ignore_ascii_case(mention) {
                return Ok(Some((
                    entity.id.clone(),
                    entity.name.clone(),
                    0.99,
                    LinkReason::LexicalMatch,
                )));
            }

            // Check semantic similarity
            let similarity = self.compute_similarity(mention, &entity.name);
            if similarity > 0.7 {
                if best.is_none() || similarity > best.as_ref().unwrap().2 {
                    best = Some((
                        entity.id.clone(),
                        entity.name.clone(),
                        similarity,
                        LinkReason::SemanticMatch,
                    ));
                }
            }

            // Check acronym (e.g., "k8s" for "Kubernetes")
            if self.is_acronym(mention, &entity.name) {
                return Ok(Some((
                    entity.id.clone(),
                    entity.name.clone(),
                    0.95,
                    LinkReason::AcronymMatch,
                )));
            }
        }

        Ok(best)
    }

    /// Fetch all entities for a project
    async fn fetch_entities(&self, project_id: &str) -> Result<Vec<EntityInfo>, String> {
        // Stub: would query database
        // For now, return empty
        Ok(vec![])
    }

    /// Compute string similarity (Jaro-Winkler style)
    fn compute_similarity(&self, s1: &str, s2: &str) -> f32 {
        let s1_lower = s1.to_lowercase();
        let s2_lower = s2.to_lowercase();

        if s1_lower == s2_lower {
            return 1.0;
        }

        if s1_lower.contains(&s2_lower) || s2_lower.contains(&s1_lower) {
            return 0.85;
        }

        // Simple Levenshtein-based similarity
        let distance = self.edit_distance(&s1_lower, &s2_lower);
        let max_len = s1_lower.len().max(s2_lower.len());
        1.0 - (distance as f32 / max_len as f32)
    }

    /// Edit distance (Levenshtein)
    fn edit_distance(&self, s1: &str, s2: &str) -> usize {
        let len1 = s1.len();
        let len2 = s2.len();
        let mut dp = vec![vec![0; len2 + 1]; len1 + 1];

        for i in 0..=len1 {
            dp[i][0] = i;
        }
        for j in 0..=len2 {
            dp[0][j] = j;
        }

        for (i, c1) in s1.chars().enumerate() {
            for (j, c2) in s2.chars().enumerate() {
                let cost = if c1 == c2 { 0 } else { 1 };
                dp[i + 1][j + 1] =
                    (dp[i][j + 1] + 1).min(dp[i + 1][j] + 1).min(dp[i][j] + cost);
            }
        }

        dp[len1][len2]
    }

    /// Check if s1 is acronym of s2
    fn is_acronym(&self, s1: &str, s2: &str) -> bool {
        if s1.len() > s2.len() || s1.is_empty() {
            return false;
        }
        let words: Vec<&str> = s2.split_whitespace().collect();
        let acronym: String = words.iter().filter_map(|w| w.chars().next()).collect();
        acronym.to_lowercase() == s1.to_lowercase()
    }

    /// Check if two strings are similar
    fn is_similar(&self, s1: &str, s2: &str) -> bool {
        self.compute_similarity(s1, s2) > 0.7
    }

    /// Check if two entities have common relations (stub)
    fn have_common_relations(&self, _id1: &str, _id2: &str) -> bool {
        // TODO: Query edge table for common neighbors
        false
    }
}

/// Internal mention structure
struct Mention {
    text: String,
    start: usize,
    end: usize,
}

/// Entity info for linking
struct EntityInfo {
    id: String,
    name: String,
}

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

    fn create_linker_mock() -> EntityLinker {
        // Create with in-memory pool (stub for testing)
        let pool = sqlx::postgres::PgPoolOptions::new()
            .max_connections(1)
            .build_lazy();
        EntityLinker::new(pool)
    }

    #[test]
    fn test_extract_mentions_basic() {
        let linker = create_linker_mock();
        let text = "Kubernetes is a container orchestration platform.";
        let mentions = linker.extract_mentions(text).unwrap();
        assert!(mentions.len() > 0);
    }

    #[test]
    fn test_extract_mentions_multiword() {
        let linker = create_linker_mock();
        let text = "Google Cloud Platform provides services.";
        let mentions = linker.extract_mentions(text).unwrap();
        assert!(mentions.iter().any(|m| m.text.contains("Cloud")));
    }

    #[test]
    fn test_mention_link_structure() {
        let link = MentionLink {
            mention_text: "Kubernetes".to_string(),
            start_offset: 0,
            end_offset: 10,
            entity_id: "e1".to_string(),
            entity_name: "Kubernetes".to_string(),
            confidence: 0.95,
            reason: LinkReason::LexicalMatch,
        };
        assert_eq!(link.confidence, 0.95);
    }

    #[test]
    fn test_link_reason_enum() {
        let reasons = vec![
            LinkReason::SemanticMatch,
            LinkReason::LexicalMatch,
            LinkReason::AliasMatch,
            LinkReason::AcronymMatch,
            LinkReason::PartialMatch,
        ];
        assert_eq!(reasons.len(), 5);
    }

    #[test]
    fn test_alias_suggestion_structure() {
        let alias = AliasSuggestion {
            entity_id: "e1".to_string(),
            canonical_name: "Kubernetes".to_string(),
            alias: "k8s".to_string(),
            confidence: 0.9,
            frequency: 5,
        };
        assert_eq!(alias.frequency, 5);
    }

    #[test]
    fn test_merge_suggestion_structure() {
        let merge = MergeSuggestion {
            entity1_id: "e1".to_string(),
            entity1_name: "Kubernetes".to_string(),
            entity2_id: "e2".to_string(),
            entity2_name: "K8s".to_string(),
            confidence: 0.85,
            reasons: vec!["Acronym match".to_string()],
        };
        assert_eq!(merge.confidence, 0.85);
        assert_eq!(merge.reasons.len(), 1);
    }

    #[test]
    fn test_coreference_cluster_structure() {
        let cluster = CoreferenceCluster {
            entity_id: "e1".to_string(),
            mentions: vec!["Kubernetes".to_string(), "k8s".to_string()],
            mention_count: 2,
            confidence: 0.85,
        };
        assert_eq!(cluster.mention_count, 2);
    }

    #[test]
    fn test_edit_distance() {
        let linker = create_linker_mock();
        let dist = linker.edit_distance("Kubernetes", "kubernetes");
        assert_eq!(dist, 0); // Same lowercase
    }

    #[test]
    fn test_edit_distance_typo() {
        let linker = create_linker_mock();
        let dist = linker.edit_distance("Kubernetes", "Kubenetes");
        assert!(dist > 0 && dist < 5);
    }

    #[test]
    fn test_compute_similarity_exact() {
        let linker = create_linker_mock();
        let sim = linker.compute_similarity("test", "test");
        assert_eq!(sim, 1.0);
    }

    #[test]
    fn test_compute_similarity_case_insensitive() {
        let linker = create_linker_mock();
        let sim = linker.compute_similarity("Test", "test");
        assert_eq!(sim, 1.0);
    }

    #[test]
    fn test_compute_similarity_substring() {
        let linker = create_linker_mock();
        let sim = linker.compute_similarity("Kubernetes", "kubernetes");
        assert!(sim > 0.8);
    }

    #[test]
    fn test_is_acronym_true() {
        let linker = create_linker_mock();
        let is_acr = linker.is_acronym("k8s", "Kubernetes");
        assert!(is_acr);
    }

    #[test]
    fn test_is_acronym_false() {
        let linker = create_linker_mock();
        let is_acr = linker.is_acronym("test", "Kubernetes");
        assert!(!is_acr);
    }

    #[test]
    fn test_is_similar_true() {
        let linker = create_linker_mock();
        let similar = linker.is_similar("Kubernetes", "kubernetes");
        assert!(similar);
    }

    #[test]
    fn test_is_similar_false() {
        let linker = create_linker_mock();
        let similar = linker.is_similar("test", "completely different");
        assert!(!similar);
    }

    #[test]
    fn test_mention_link_reason_serialization() {
        let reason = LinkReason::SemanticMatch;
        let json = serde_json::to_string(&reason).unwrap();
        assert!(json.contains("SemanticMatch"));
    }

    #[test]
    fn test_mention_link_full_serialization() {
        let link = MentionLink {
            mention_text: "Kubernetes".to_string(),
            start_offset: 0,
            end_offset: 10,
            entity_id: "e1".to_string(),
            entity_name: "Kubernetes".to_string(),
            confidence: 0.95,
            reason: LinkReason::LexicalMatch,
        };
        let json = serde_json::to_string(&link).unwrap();
        assert!(json.contains("Kubernetes"));
        assert!(json.contains("0.95"));
    }
}