cont/src/search.rs

use std::{cmp::Reverse, path::Path};

use crate::session::{Agent, Session};
use nucleo_matcher::{
    Config, Matcher, Utf32Str,
    pattern::{CaseMatching, Normalization, Pattern},
};

pub struct SearchIndex {
    pub sessions: Vec<Session>,
    text: Vec<SearchText>,
}
struct SearchText {
    title: String,
    messages: Vec<String>,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MatchKind {
    Recent,
    Title,
    Message,
}
#[derive(Clone, Debug)]
pub struct Hit {
    pub index: usize,
    pub kind: MatchKind,
    pub message: Option<usize>,
    score: u32,
    tier: u8,
}

impl SearchIndex {
    pub fn new(sessions: Vec<Session>) -> Self {
        let text = sessions
            .iter()
            .map(|s| SearchText {
                title: s.title.to_lowercase(),
                messages: s.messages.iter().map(|s| s.to_lowercase()).collect(),
            })
            .collect();
        Self { sessions, text }
    }

    pub fn search(&self, query: &str, agent: Option<Agent>, cwd: Option<&Path>) -> Vec<Hit> {
        let query = query.trim().to_lowercase();
        let words: Vec<_> = query.split_whitespace().collect();
        let pattern = Pattern::new(
            &query,
            CaseMatching::Ignore,
            Normalization::Smart,
            nucleo_matcher::pattern::AtomKind::Fuzzy,
        );
        let mut matcher = Matcher::new(Config::DEFAULT);
        let mut buffer = vec![];
        let mut hits = vec![];
        for (index, (session, text)) in self.sessions.iter().zip(&self.text).enumerate() {
            if agent.is_some_and(|a| a != session.agent) || cwd.is_some_and(|p| p != session.cwd) {
                continue;
            }
            let mut hit = Hit {
                index,
                kind: MatchKind::Recent,
                message: None,
                score: 0,
                tier: 0,
            };
            if !query.is_empty() {
                if text.title.contains(&query) {
                    hit.kind = MatchKind::Title;
                    hit.tier = 4;
                    hit.score = if text.title == query {
                        2
                    } else if text.title.starts_with(&query) {
                        1
                    } else {
                        0
                    };
                } else if words.iter().all(|w| text.title.contains(w)) {
                    hit.kind = MatchKind::Title;
                    hit.tier = 3;
                } else if let Some(score) =
                    pattern.score(Utf32Str::new(&session.title, &mut buffer), &mut matcher)
                {
                    hit.kind = MatchKind::Title;
                    hit.tier = 2;
                    hit.score = score;
                } else {
                    let mut found = None;
                    for (i, message) in text.messages.iter().enumerate() {
                        if message.contains(&query) {
                            found = Some((i, 1));
                            break;
                        }
                        if found.is_none() && words.iter().all(|w| message.contains(w)) {
                            found = Some((i, 0));
                        }
                    }
                    let Some((i, score)) = found else { continue };
                    hit.kind = MatchKind::Message;
                    hit.tier = 1;
                    hit.score = score;
                    hit.message = Some(i);
                }
            }
            hits.push(hit);
        }
        hits.sort_by_key(|h| {
            (
                Reverse(h.tier),
                Reverse(h.score),
                Reverse(self.sessions[h.index].modified),
                h.index,
            )
        });
        hits
    }
}

/// Context around the matching portion, not just the start of a long prompt.
pub fn snippet(text: &str, query: &str, max_chars: usize) -> String {
    let lower = text.to_lowercase();
    let query = query.to_lowercase();
    let pos = query
        .split_whitespace()
        .filter_map(|w| lower.find(w))
        .min()
        .unwrap_or(0);
    // Lowercasing can change byte lengths. Convert to a character offset first.
    let char_pos = lower[..pos].chars().count();
    let start = char_pos.saturating_sub(70);
    let mut output = if start > 0 {
        "…".to_owned()
    } else {
        String::new()
    };
    output.extend(text.chars().skip(start).take(max_chars));
    if text.chars().count() > start + max_chars {
        output.push('…');
    }
    output
}

#[cfg(test)]
mod tests {
    use super::*;
    fn session(title: &str, message: &str, modified: u64) -> Session {
        Session {
            agent: Agent::Pi,
            id: title.into(),
            path: "/x".into(),
            cwd: "/work".into(),
            title: title.into(),
            modified,
            messages: vec![message.into()],
        }
    }
    #[test]
    fn title_hits_always_beat_body_hits_and_recency_breaks_ties() {
        let index = SearchIndex::new(vec![
            session("Other", "authentication failure", 999),
            session("Authentication", "", 1),
            session("Auth guide", "", 2),
        ]);
        let hits = index.search("auth", None, None);
        assert_eq!(hits.iter().map(|h| h.index).collect::<Vec<_>>(), [2, 1, 0]);
        assert_eq!(hits[2].kind, MatchKind::Message);
    }
    #[test]
    fn fuzzy_titles_and_unordered_message_tokens() {
        let index = SearchIndex::new(vec![
            session("Authentication refactor", "", 1),
            session("Other", "cache invalidation fix", 2),
        ]);
        assert_eq!(index.search("athref", None, None)[0].index, 0);
        let hits = index.search("fix cache", None, None);
        assert_eq!(hits[0].index, 1);
        assert_eq!(hits[0].message, Some(0));
        assert!(index.search("no match at all", None, None).is_empty());
    }
    #[test]
    fn scope_is_exact_and_filters_compose() {
        let index = SearchIndex::new(vec![session("title", "text", 1)]);
        assert!(index.search("", Some(Agent::Claude), None).is_empty());
        assert!(index.search("", None, Some(Path::new("/wor"))).is_empty());
        assert_eq!(
            index
                .search("", Some(Agent::Pi), Some(Path::new("/work")))
                .len(),
            1
        );
    }
    #[test]
    fn unicode_and_match_context() {
        let index = SearchIndex::new(vec![session("École 日本語", "the café is open", 1)]);
        assert_eq!(index.search("日本", None, None).len(), 1);
        assert_eq!(index.search("CAFÉ", None, None).len(), 1);
        let text = format!("{}needle", "x".repeat(1000));
        assert!(snippet(&text, "needle", 100).contains("needle"));
    }
}