goosedump 0.5.1

// SPDX-License-Identifier: LGPL-2.1-or-later
// Copyright (C) Jarkko Sakkinen 2026

use crate::display;
use crate::message::{ConversationMessage, SearchHit};
use crate::text;
use std::collections::{HashMap, HashSet};

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

#[allow(clippy::cast_precision_loss)]
fn usize_as_f64(value: usize) -> f64 {
    value as f64
}

fn highlight_terms(query: &str) -> Vec<String> {
    text::split_words(query)
        .into_iter()
        .filter(|word| word.chars().count() > 1)
        .collect()
}

fn min_span_all(positions: &[&[usize]]) -> usize {
    let mut all: Vec<(usize, usize)> = Vec::new();
    for (term_idx, pos_list) in positions.iter().enumerate() {
        for &pos in *pos_list {
            all.push((pos, term_idx));
        }
    }
    all.sort_by_key(|(pos, _)| *pos);

    let num_terms = positions.len();
    let mut counts = vec![0usize; num_terms];
    let mut matched = 0;
    let mut left = 0;
    let mut min_span = usize::MAX;

    for right in 0..all.len() {
        let term_idx = all[right].1;
        if counts[term_idx] == 0 {
            matched += 1;
        }
        counts[term_idx] += 1;

        while matched == num_terms {
            let span = all[right].0 - all[left].0 + 1;
            min_span = min_span.min(span);

            let left_term = all[left].1;
            counts[left_term] -= 1;
            if counts[left_term] == 0 {
                matched -= 1;
            }
            left += 1;
        }
    }

    min_span
}

fn adjacent_fallback(
    messages: &[ConversationMessage],
    tokenized: &[Vec<String>],
    terms: &[String],
    page: usize,
    page_size: usize,
) -> (Vec<SearchHit>, usize) {
    let hits: Vec<SearchHit> = tokenized
        .iter()
        .enumerate()
        .filter_map(|(i, msg_words)| {
            let word_set: HashSet<&str> =
                msg_words.iter().map(std::string::String::as_str).collect();
            let matching_pairs = terms
                .windows(2)
                .filter(|pair| {
                    word_set.contains(pair[0].as_str()) && word_set.contains(pair[1].as_str())
                })
                .count();
            if matching_pairs == 0 {
                return None;
            }
            let msg = &messages[i];
            let role = msg.role_label();
            let searchable = display::searchable_text(msg);
            let snippet = text::line_snippet_terms(&searchable, terms, 2);
            let files = display::message_files(msg);
            Some(SearchHit {
                entry_id: msg.entry_id.clone(),
                score: usize_as_f64(matching_pairs),
                role,
                text: snippet,
                files,
                terms: terms.to_vec(),
            })
        })
        .collect();

    sort_and_page(hits, page, page_size)
}

fn sort_and_page(
    mut hits: Vec<SearchHit>,
    page: usize,
    page_size: usize,
) -> (Vec<SearchHit>, usize) {
    sort_by_score_desc(&mut hits);
    let total = hits.len();
    let start = (page.saturating_sub(1)) * page_size;
    if start >= total {
        return (Vec::new(), total);
    }
    (
        hits[start..std::cmp::min(start + page_size, total)].to_vec(),
        total,
    )
}

fn sort_by_score_desc(hits: &mut [SearchHit]) {
    hits.sort_by(|a, b| {
        b.score
            .partial_cmp(&a.score)
            .unwrap_or(std::cmp::Ordering::Equal)
    });
}

fn grep_hit(msg: &ConversationMessage, pattern_lc: &str, pattern: &str) -> Option<SearchHit> {
    let searchable = display::searchable_text(msg);
    let searchable_lc = searchable.to_ascii_lowercase();
    if !text::glob_search(pattern_lc, &searchable_lc) {
        return None;
    }
    let match_lines = searchable_lc
        .split('\n')
        .filter(|line| text::glob_search(pattern_lc, line))
        .count()
        .max(1);
    Some(SearchHit {
        entry_id: msg.entry_id.clone(),
        score: usize_as_f64(match_lines),
        role: msg.role_label(),
        text: text::line_snippet_glob(&searchable, pattern_lc, 2),
        files: display::message_files(msg),
        terms: highlight_terms(pattern),
    })
}

/// Filter messages whose searchable text matches the glob `pattern`,
/// case-insensitively, preserving transcript order.
pub fn grep(messages: &[ConversationMessage], pattern: &str) -> Vec<SearchHit> {
    let pattern_lc = pattern.to_ascii_lowercase();
    messages
        .iter()
        .filter_map(|msg| grep_hit(msg, &pattern_lc, pattern))
        .collect()
}

struct TermSearchIndex {
    tokenized: Vec<Vec<String>>,
    doc_lengths: Vec<usize>,
    df: HashMap<String, usize>,
    avgdl: f64,
    n: f64,
}

fn normalized_terms(query_str: &str) -> Vec<String> {
    text::split_words(query_str)
        .into_iter()
        .filter(|w| w.len() > 1 && !text::is_stop_word(w))
        .collect()
}

fn build_term_search_index(messages: &[ConversationMessage], terms: &[String]) -> TermSearchIndex {
    let n = usize_as_f64(messages.len());
    let mut df: HashMap<String, usize> = HashMap::new();
    let mut tokenized: Vec<Vec<String>> = Vec::with_capacity(messages.len());
    let mut doc_lengths: Vec<usize> = Vec::with_capacity(messages.len());
    let mut total_words: usize = 0;

    for msg in messages {
        let searchable = display::searchable_text(msg);
        let msg_words: Vec<String> = text::split_words(&searchable)
            .into_iter()
            .map(|w| w.to_ascii_lowercase())
            .collect();
        let dl = msg_words.len();
        doc_lengths.push(dl);
        total_words += dl;

        let word_set: HashSet<&str> = msg_words.iter().map(String::as_str).collect();
        for term in terms {
            if word_set.contains(term.as_str()) {
                *df.entry(term.clone()).or_insert(0) += 1;
            }
        }
        tokenized.push(msg_words);
    }

    let avgdl = if n > 0.0 {
        usize_as_f64(total_words) / n
    } else {
        0.0
    };

    TermSearchIndex {
        tokenized,
        doc_lengths,
        df,
        avgdl,
        n,
    }
}

fn positions_by_term<'a>(
    msg_words: &'a [String],
    term_set: &HashSet<&str>,
) -> HashMap<&'a str, Vec<usize>> {
    let mut positions_by_term: HashMap<&str, Vec<usize>> = HashMap::new();
    for (pos, word) in msg_words.iter().enumerate() {
        let word = word.as_str();
        if term_set.contains(word) {
            positions_by_term.entry(word).or_default().push(pos);
        }
    }
    positions_by_term
}

fn score_message_terms(
    index: &TermSearchIndex,
    terms: &[String],
    positions_by_term: &HashMap<&str, Vec<usize>>,
    doc_index: usize,
) -> f64 {
    let dl = usize_as_f64(index.doc_lengths[doc_index]);
    let len_norm = 1.0 - BM25_B + BM25_B * (dl / index.avgdl.max(1.0));
    let mut score = 0.0;
    let mut term_positions: Vec<&[usize]> = Vec::with_capacity(terms.len());

    for term in terms {
        let Some(positions) = positions_by_term.get(term.as_str()) else {
            continue;
        };
        let tf = positions.len();
        term_positions.push(positions.as_slice());
        let df_val = usize_as_f64(*index.df.get(term).unwrap_or(&1));
        let idf = (1.0 + (index.n - df_val + 0.5) / (df_val + 0.5)).ln();
        let tf = usize_as_f64(tf);
        let tf_score = (tf * (BM25_K1 + 1.0)) / (tf + BM25_K1 * len_norm);
        score += idf * tf_score;
    }

    if score > 0.0 && term_positions.len() > 1 {
        let min_span = min_span_all(&term_positions).max(1);
        score *= 1.0 + PROXIMITY_FACTOR / (1.0 + usize_as_f64(min_span));
    }

    score
}

fn term_search_hit(msg: &ConversationMessage, terms: &[String], score: f64) -> SearchHit {
    let role = msg.role_label();
    let searchable = display::searchable_text(msg);
    let snippet = text::line_snippet_terms(&searchable, terms, 2);
    let files = display::message_files(msg);

    SearchHit {
        entry_id: msg.entry_id.clone(),
        score,
        role,
        text: snippet,
        files,
        terms: terms.to_vec(),
    }
}

fn query_terms(
    messages: &[ConversationMessage],
    terms: &[String],
    page: usize,
    page_size: usize,
) -> (Vec<SearchHit>, usize) {
    let index = build_term_search_index(messages, terms);
    let term_set: HashSet<&str> = terms.iter().map(String::as_str).collect();
    let hits: Vec<SearchHit> = messages
        .iter()
        .enumerate()
        .filter_map(|(i, msg)| {
            let positions_by_term = positions_by_term(&index.tokenized[i], &term_set);
            let score = score_message_terms(&index, terms, &positions_by_term, i);
            (score > 0.0).then(|| term_search_hit(msg, terms, score))
        })
        .collect();

    let (hits, total) = sort_and_page(hits, page, page_size);
    if hits.is_empty() && total == 0 && terms.len() >= 2 {
        return adjacent_fallback(messages, &index.tokenized, terms, page, page_size);
    }
    (hits, total)
}

pub fn query(
    messages: &[ConversationMessage],
    query_str: &str,
    page: usize,
    page_size: usize,
) -> (Vec<SearchHit>, usize) {
    let terms = normalized_terms(query_str);
    if terms.is_empty() {
        return (Vec::new(), 0);
    }

    query_terms(messages, &terms, page, page_size)
}