kg/

output.rs

use std::cmp::Reverse;

use nucleo_matcher::pattern::{CaseMatching, Normalization, Pattern};
use nucleo_matcher::{Config, Matcher, Utf32Str};

use crate::graph::{Edge, GraphFile, Node};
use crate::index::Bm25Index;

const BM25_K1: f64 = 1.5;
const BM25_B: f64 = 0.75;

#[derive(Debug, Clone, Copy)]
pub enum FindMode {
    Fuzzy,
    Bm25,
}

pub fn render_find(
    graph: &GraphFile,
    queries: &[String],
    limit: usize,
    include_features: bool,
    mode: FindMode,
    full: bool,
) -> String {
    render_find_with_index(graph, queries, limit, include_features, mode, full, None)
}

pub fn render_find_with_index(
    graph: &GraphFile,
    queries: &[String],
    limit: usize,
    include_features: bool,
    mode: FindMode,
    full: bool,
    index: Option<&Bm25Index>,
) -> String {
    let mut sections = Vec::new();
    for query in queries {
        let matches = find_matches_with_index(graph, query, limit, include_features, mode, index);
        let mut lines = vec![format!("? {query} ({})", matches.len())];
        for node in matches {
            lines.push(render_node_block(graph, node, full));
        }
        sections.push(lines.join("\n"));
    }
    format!("{}\n", sections.join("\n\n"))
}

pub fn find_nodes(
    graph: &GraphFile,
    query: &str,
    limit: usize,
    include_features: bool,
    mode: FindMode,
) -> Vec<Node> {
    find_matches_with_index(graph, query, limit, include_features, mode, None)
        .into_iter()
        .cloned()
        .collect()
}

pub fn find_nodes_with_index(
    graph: &GraphFile,
    query: &str,
    limit: usize,
    include_features: bool,
    mode: FindMode,
    index: Option<&Bm25Index>,
) -> Vec<Node> {
    find_matches_with_index(graph, query, limit, include_features, mode, index)
        .into_iter()
        .cloned()
        .collect()
}

pub fn count_find_results(
    graph: &GraphFile,
    queries: &[String],
    limit: usize,
    include_features: bool,
    mode: FindMode,
) -> usize {
    count_find_results_with_index(graph, queries, limit, include_features, mode, None)
}

pub fn count_find_results_with_index(
    graph: &GraphFile,
    queries: &[String],
    limit: usize,
    include_features: bool,
    mode: FindMode,
    index: Option<&Bm25Index>,
) -> usize {
    let mut total = 0;
    for query in queries {
        let matches = find_matches_with_index(graph, query, limit, include_features, mode, index);
        total += matches.len();
    }
    total
}

pub fn render_node(graph: &GraphFile, node: &Node, full: bool) -> String {
    format!("{}\n", render_node_block(graph, node, full))
}

fn render_node_block(graph: &GraphFile, node: &Node, full: bool) -> String {
    let mut lines = Vec::new();
    lines.push(format!("# {} | {}", node.id, node.name));

    if !node.properties.alias.is_empty() {
        lines.push(format!("aka: {}", node.properties.alias.join(", ")));
    }
    if full {
        if !node.properties.domain_area.is_empty() {
            lines.push(format!("domain_area: {}", node.properties.domain_area));
        }
        if !node.properties.provenance.is_empty() {
            lines.push(format!("provenance: {}", node.properties.provenance));
        }
        if let Some(confidence) = node.properties.confidence {
            lines.push(format!("confidence: {confidence}"));
        }
        lines.push(format!("importance: {}", node.properties.importance));
        if !node.properties.created_at.is_empty() {
            lines.push(format!("created_at: {}", node.properties.created_at));
        }
    }

    let facts_to_show = if full {
        node.properties.key_facts.len()
    } else {
        node.properties.key_facts.len().min(2)
    };
    for fact in node.properties.key_facts.iter().take(facts_to_show) {
        lines.push(format!("- {fact}"));
    }
    if node.properties.key_facts.len() > facts_to_show || full {
        lines.push(format!("({} facts total)", node.properties.key_facts.len()));
    }

    let note_count = graph
        .notes
        .iter()
        .filter(|note| note.node_id == node.id)
        .count();
    if full && note_count > 0 {
        lines.push(format!("notes: {note_count}"));
    }

    for edge in outgoing_edges(graph, &node.id, full) {
        if let Some(target) = graph.node_by_id(&edge.target_id) {
            lines.push(format_edge("->", edge, target));
        }
    }
    for edge in incoming_edges(graph, &node.id, full) {
        if let Some(source) = graph.node_by_id(&edge.source_id) {
            lines.push(format_edge("<-", edge, source));
        }
    }

    lines.join("\n")
}

fn outgoing_edges<'a>(graph: &'a GraphFile, node_id: &str, full: bool) -> Vec<&'a Edge> {
    let mut edges: Vec<&Edge> = graph
        .edges
        .iter()
        .filter(|edge| edge.source_id == node_id)
        .collect();
    edges.sort_by_key(|edge| (&edge.relation, &edge.target_id));
    if !full {
        edges.truncate(3);
    }
    edges
}

fn incoming_edges<'a>(graph: &'a GraphFile, node_id: &str, full: bool) -> Vec<&'a Edge> {
    let mut edges: Vec<&Edge> = graph
        .edges
        .iter()
        .filter(|edge| edge.target_id == node_id)
        .collect();
    edges.sort_by_key(|edge| (&edge.relation, &edge.source_id));
    if !full {
        edges.truncate(3);
    }
    edges
}

fn format_edge(prefix: &str, edge: &Edge, related: &Node) -> String {
    let (arrow, relation) = if edge.relation.starts_with("NOT_") {
        (
            format!("{prefix}!"),
            edge.relation.trim_start_matches("NOT_"),
        )
    } else {
        (prefix.to_owned(), edge.relation.as_str())
    };

    let mut line = format!("{arrow} {relation} | {} | {}", related.id, related.name);
    if !edge.properties.detail.is_empty() {
        line.push_str(" | ");
        line.push_str(&truncate(&edge.properties.detail, 80));
    }
    line
}

fn truncate(value: &str, max_len: usize) -> String {
    let char_count = value.chars().count();
    if char_count <= max_len {
        return value.to_owned();
    }
    let truncated: String = value.chars().take(max_len.saturating_sub(3)).collect();
    format!("{truncated}...")
}

fn find_matches_with_index<'a>(
    graph: &'a GraphFile,
    query: &str,
    limit: usize,
    include_features: bool,
    mode: FindMode,
    index: Option<&Bm25Index>,
) -> Vec<&'a Node> {
    let mut scored: Vec<(i64, Reverse<&str>, &'a Node)> = match mode {
        FindMode::Fuzzy => {
            let pattern = Pattern::parse(query, CaseMatching::Ignore, Normalization::Smart);
            let mut matcher = Matcher::new(Config::DEFAULT);
            graph
                .nodes
                .iter()
                .filter(|node| include_features || node.r#type != "Feature")
                .filter_map(|node| {
                    score_node(node, query, &pattern, &mut matcher).map(|score| {
                        let base = score as i64;
                        let boost = feedback_boost(node);
                        (base + boost, Reverse(node.id.as_str()), node)
                    })
                })
                .collect()
        }
        FindMode::Bm25 => score_bm25(graph, query, include_features, index),
    };

    scored.sort_by(|left, right| right.0.cmp(&left.0).then_with(|| left.1.cmp(&right.1)));
    scored
        .into_iter()
        .take(limit)
        .map(|(_, _, node)| node)
        .collect()
}

fn feedback_boost(node: &Node) -> i64 {
    let count = node.properties.feedback_count as f64;
    if count <= 0.0 {
        return 0;
    }
    let avg = node.properties.feedback_score / count;
    let confidence = (count.ln_1p() / 3.0).min(1.0);
    let scaled = avg * 200.0 * confidence;
    scaled.clamp(-300.0, 300.0).round() as i64
}

fn score_bm25<'a>(
    graph: &'a GraphFile,
    query: &str,
    include_features: bool,
    index: Option<&Bm25Index>,
) -> Vec<(i64, Reverse<&'a str>, &'a Node)> {
    let terms = tokenize(query);
    if terms.is_empty() {
        return Vec::new();
    }

    if let Some(idx) = index {
        let results = idx.search(&terms, graph);
        return results
            .into_iter()
            .filter_map(|(node_id, score)| {
                let node = graph.node_by_id(&node_id)?;
                if !include_features && node.r#type == "Feature" {
                    return None;
                }
                let boost = feedback_boost(node) as f64;
                let combined = (score as f64 * 100.0 + boost).round() as i64;
                Some((combined, Reverse(node.id.as_str()), node))
            })
            .collect();
    }

    let mut docs: Vec<(&'a Node, Vec<String>)> = graph
        .nodes
        .iter()
        .filter(|node| include_features || node.r#type != "Feature")
        .map(|node| (node, tokenize(&node_document_text(graph, node))))
        .collect();

    if docs.is_empty() {
        return Vec::new();
    }

    let mut df: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
    for term in &terms {
        let mut count = 0usize;
        for (_, tokens) in &docs {
            if tokens.iter().any(|t| t == term) {
                count += 1;
            }
        }
        df.insert(term.as_str(), count);
    }

    let total_docs = docs.len() as f64;
    let avgdl = docs
        .iter()
        .map(|(_, tokens)| tokens.len() as f64)
        .sum::<f64>()
        / total_docs;

    let mut scored = Vec::new();

    for (node, tokens) in docs.drain(..) {
        let dl = tokens.len() as f64;
        if dl == 0.0 {
            continue;
        }
        let mut score = 0.0f64;
        for term in &terms {
            let tf = tokens.iter().filter(|t| *t == term).count() as f64;
            if tf == 0.0 {
                continue;
            }
            let df_t = *df.get(term.as_str()).unwrap_or(&0) as f64;
            let idf = (1.0 + (total_docs - df_t + 0.5) / (df_t + 0.5)).ln();
            let denom = tf + BM25_K1 * (1.0 - BM25_B + BM25_B * (dl / avgdl));
            score += idf * (tf * (BM25_K1 + 1.0) / denom);
        }
        if score > 0.0 {
            let boost = feedback_boost(node) as f64;
            let combined = score * 100.0 + boost;
            scored.push((combined.round() as i64, Reverse(node.id.as_str()), node));
        }
    }

    scored
}

fn node_document_text(graph: &GraphFile, node: &Node) -> String {
    let mut out = String::new();
    push_field(&mut out, &node.id);
    push_field(&mut out, &node.name);
    push_field(&mut out, &node.properties.description);
    for alias in &node.properties.alias {
        push_field(&mut out, alias);
    }
    for fact in &node.properties.key_facts {
        push_field(&mut out, fact);
    }
    for note in graph.notes.iter().filter(|note| note.node_id == node.id) {
        push_field(&mut out, &note.body);
        for tag in &note.tags {
            push_field(&mut out, tag);
        }
    }
    out
}

fn push_field(target: &mut String, value: &str) {
    if value.is_empty() {
        return;
    }
    if !target.is_empty() {
        target.push(' ');
    }
    target.push_str(value);
}

fn tokenize(text: &str) -> Vec<String> {
    let mut tokens = Vec::new();
    let mut current = String::new();
    for ch in text.chars() {
        if ch.is_alphanumeric() {
            current.push(ch.to_ascii_lowercase());
        } else if !current.is_empty() {
            tokens.push(current.clone());
            current.clear();
        }
    }
    if !current.is_empty() {
        tokens.push(current);
    }
    tokens
}

fn score_node(node: &Node, query: &str, pattern: &Pattern, matcher: &mut Matcher) -> Option<u32> {
    let mut total = 0;
    let mut primary_hits = 0;

    let id_score = score_primary_field(query, pattern, matcher, &node.id, 4);
    if id_score > 0 {
        primary_hits += 1;
    }
    total += id_score;

    let name_score = score_primary_field(query, pattern, matcher, &node.name, 3);
    if name_score > 0 {
        primary_hits += 1;
    }
    total += name_score;

    for alias in &node.properties.alias {
        let alias_score = score_primary_field(query, pattern, matcher, alias, 3);
        if alias_score > 0 {
            primary_hits += 1;
        }
        total += alias_score;
    }

    if primary_hits > 0 {
        total += score_secondary_field(query, pattern, matcher, &node.properties.description, 1);
    }

    (total > 0).then_some(total)
}

fn score_field(pattern: &Pattern, matcher: &mut Matcher, value: &str) -> Option<u32> {
    if value.is_empty() {
        return None;
    }
    let mut buf = Vec::new();
    let haystack = Utf32Str::new(value, &mut buf);
    pattern.score(haystack, matcher)
}

fn score_primary_field(
    query: &str,
    pattern: &Pattern,
    matcher: &mut Matcher,
    value: &str,
    weight: u32,
) -> u32 {
    let bonus = textual_bonus(query, value);
    if bonus == 0 {
        return 0;
    }
    let fuzzy = score_field(pattern, matcher, value).unwrap_or(0);
    (fuzzy + bonus) * weight
}

fn score_secondary_field(
    query: &str,
    pattern: &Pattern,
    matcher: &mut Matcher,
    value: &str,
    weight: u32,
) -> u32 {
    let bonus = textual_bonus(query, value);
    if bonus == 0 {
        return 0;
    }
    let fuzzy = score_field(pattern, matcher, value).unwrap_or(0);
    (fuzzy + bonus / 2) * weight
}

fn textual_bonus(query: &str, value: &str) -> u32 {
    let query = query.trim().to_lowercase();
    let value = value.to_lowercase();

    if value == query {
        return 400;
    }
    if value.contains(&query) {
        return 200;
    }

    query
        .split_whitespace()
        .map(|token| {
            if value.contains(token) {
                80
            } else if is_subsequence(token, &value) {
                40
            } else {
                0
            }
        })
        .sum()
}

fn is_subsequence(needle: &str, haystack: &str) -> bool {
    if needle.is_empty() {
        return false;
    }

    let mut chars = needle.chars();
    let mut current = match chars.next() {
        Some(ch) => ch,
        None => return false,
    };

    for ch in haystack.chars() {
        if ch == current {
            match chars.next() {
                Some(next) => current = next,
                None => return true,
            }
        }
    }

    false
}