archlens/crates/domain/src/aggregates/code_graph.rs

use std::collections::{HashMap, HashSet};

use crate::{CodeElement, ModuleName, Relationship, RelationshipKind};

#[derive(Debug, Clone)]
pub struct CodeGraph {
    elements: Vec<CodeElement>,
    relationships: Vec<Relationship>,
}

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

impl CodeGraph {
    pub fn new() -> Self {
        Self {
            elements: Vec::new(),
            relationships: Vec::new(),
        }
    }

    pub fn add_element(&mut self, element: CodeElement) {
        self.elements.push(element);
    }

    pub fn add_relationship(&mut self, relationship: Relationship) {
        self.relationships.push(relationship);
    }

    pub fn elements(&self) -> &[CodeElement] {
        &self.elements
    }

    pub fn relationships(&self) -> &[Relationship] {
        &self.relationships
    }

    pub fn modules(&self) -> Vec<ModuleName> {
        let mut seen = HashSet::new();
        let mut modules = Vec::new();

        for element in &self.elements {
            if let Some(module) = element.module()
                && seen.insert(module.as_str().to_string())
            {
                modules.push(module.clone());
            }
        }

        modules
    }

    pub fn elements_by_module(&self) -> (HashMap<String, Vec<&CodeElement>>, Vec<&CodeElement>) {
        let mut grouped: HashMap<String, Vec<&CodeElement>> = HashMap::new();
        let mut ungrouped: Vec<&CodeElement> = Vec::new();

        for element in &self.elements {
            if let Some(module) = element.module() {
                grouped
                    .entry(module.as_str().to_string())
                    .or_default()
                    .push(element);
            } else {
                ungrouped.push(element);
            }
        }

        (grouped, ungrouped)
    }

    pub fn resolve_relationships(self) -> CodeGraph {
        let qualified_names: HashSet<&str> =
            self.elements.iter().map(|e| e.qualified_name()).collect();

        // Also keep bare name lookup for import relationships and unqualified fallback
        let all_bare_names: HashSet<&str> = self.elements.iter().map(|e| e.name()).collect();

        let mut resolved = CodeGraph::new();
        for element in &self.elements {
            resolved.add_element(element.clone());
        }
        for rel in &self.relationships {
            match rel.kind() {
                RelationshipKind::Import => {
                    resolved.add_relationship(rel.clone());
                }
                _ => {
                    let src_ok = qualified_names.contains(rel.source())
                        || all_bare_names.contains(rel.source());
                    let tgt_ok = qualified_names.contains(rel.target())
                        || all_bare_names.contains(rel.target());
                    if src_ok && tgt_ok {
                        resolved.add_relationship(rel.clone());
                    }
                }
            }
        }
        resolved
    }

    pub fn filter_external_imports(self, known_modules: &HashSet<String>) -> CodeGraph {
        let module_names: HashSet<String> = self
            .modules()
            .iter()
            .map(|m| m.as_str().to_lowercase())
            .collect();

        let all_known: HashSet<&str> = known_modules
            .iter()
            .map(|s| s.as_str())
            .chain(module_names.iter().map(|s| s.as_str()))
            .collect();

        let mut filtered = CodeGraph::new();
        for element in &self.elements {
            filtered.add_element(element.clone());
        }
        for rel in &self.relationships {
            if rel.kind() == RelationshipKind::Import {
                let target_top = rel.target().split('.').next().unwrap_or("").to_lowercase();
                if !all_known.contains(target_top.as_str()) {
                    continue;
                }
            }
            filtered.add_relationship(rel.clone());
        }
        filtered
    }

    pub fn qualify(self) -> CodeGraph {
        // Build lookup: bare name -> Vec<qualified_name> (for disambiguation)
        let mut name_to_qualified: HashMap<&str, Vec<String>> = HashMap::new();
        for element in &self.elements {
            let qn = match element.module() {
                Some(m) => format!("{}::{}", m.as_str(), element.name()),
                None => element.name().to_string(),
            };
            name_to_qualified
                .entry(element.name())
                .or_default()
                .push(qn);
        }

        // Build lookup: file_path -> qualified source names in that file
        let mut file_to_qualified: HashMap<&str, Vec<String>> = HashMap::new();
        for element in &self.elements {
            let qn = match element.module() {
                Some(m) => format!("{}::{}", m.as_str(), element.name()),
                None => element.name().to_string(),
            };
            file_to_qualified
                .entry(element.file_path().as_str())
                .or_default()
                .push(qn);
        }

        let mut qualified = CodeGraph::new();

        // 1. Qualify element names
        for element in &self.elements {
            let qn = match element.module() {
                Some(m) => format!("{}::{}", m.as_str(), element.name()),
                None => element.name().to_string(),
            };
            qualified.add_element(element.clone().with_qualified_name(qn));
        }

        // 2. Rewrite relationship source/target
        for rel in &self.relationships {
            // Qualify source: find the qualified name of the source in its file
            let src_qualified = rel
                .source_file()
                .and_then(|f| file_to_qualified.get(f.as_str()))
                .and_then(|qns| {
                    qns.iter().find(|qn| {
                        qn.ends_with(&format!("::{}", rel.source())) || *qn == rel.source()
                    })
                })
                .cloned()
                .unwrap_or_else(|| {
                    // Fall back: unambiguous lookup
                    name_to_qualified
                        .get(rel.source())
                        .filter(|v| v.len() == 1)
                        .and_then(|v| v.first())
                        .cloned()
                        .unwrap_or_else(|| rel.source().to_string())
                });

            // Qualify target: prefer same module as source when ambiguous
            let src_module = src_qualified.split("::").next().unwrap_or("");
            let tgt_qualified = match name_to_qualified.get(rel.target()) {
                Some(candidates) if candidates.len() == 1 => candidates[0].clone(),
                Some(candidates) => {
                    // Prefer same module as source
                    candidates
                        .iter()
                        .find(|qn| qn.starts_with(&format!("{}::", src_module)))
                        .cloned()
                        .unwrap_or_else(|| rel.target().to_string())
                }
                None => rel.target().to_string(),
            };

            let new_rel = Relationship::new(&src_qualified, &tgt_qualified, rel.kind())
                .unwrap_or_else(|_| rel.clone());
            let new_rel = if let Some(f) = rel.source_file() {
                new_rel.with_source_file(f.clone())
            } else {
                new_rel
            };
            qualified.add_relationship(new_rel);
        }

        qualified
    }

    pub fn cross_module_deps_for(&self, module: &ModuleName) -> Vec<(ModuleName, usize)> {
        let module_element_qnames: HashSet<&str> = self
            .elements
            .iter()
            .filter(|e| e.module().is_some_and(|m| m == module))
            .map(|e| e.qualified_name())
            .collect();

        let target_module_of: HashMap<&str, Option<&ModuleName>> = self
            .elements
            .iter()
            .map(|e| (e.qualified_name(), e.module()))
            .collect();

        let mut counts: HashMap<&str, usize> = HashMap::new();
        for rel in &self.relationships {
            if !module_element_qnames.contains(rel.source()) {
                continue;
            }
            if module_element_qnames.contains(rel.target()) {
                continue;
            }
            if let Some(Some(target_mod)) = target_module_of.get(rel.target()) {
                *counts.entry(target_mod.as_str()).or_insert(0) += 1;
            }
        }

        let mut result: Vec<(ModuleName, usize)> = counts
            .into_iter()
            .filter_map(|(name, count)| ModuleName::new(name).ok().map(|m| (m, count)))
            .collect();
        result.sort_by(|a, b| a.0.as_str().cmp(b.0.as_str()));
        result
    }

    pub fn subgraph_by_module(&self, module: &ModuleName) -> CodeGraph {
        let filtered_elements: Vec<CodeElement> = self
            .elements
            .iter()
            .filter(|e| e.module().is_some_and(|m| m == module))
            .cloned()
            .collect();

        let element_qnames: HashSet<&str> = filtered_elements
            .iter()
            .map(|e| e.qualified_name())
            .collect();

        let filtered_relationships: Vec<Relationship> = self
            .relationships
            .iter()
            .filter(|r| element_qnames.contains(r.source()) && element_qnames.contains(r.target()))
            .cloned()
            .collect();

        CodeGraph {
            elements: filtered_elements,
            relationships: filtered_relationships,
        }
    }
}