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,
}