frxd/src/index.rs

use std::collections::HashSet;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::Mutex;

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use tantivy::collector::{Count, TopDocs};
use tantivy::directory::MmapDirectory;
use tantivy::query::{
    AllQuery, BooleanQuery, BoostQuery, EmptyQuery, Occur, PhraseQuery, Query, TermQuery,
};
use tantivy::schema::{
    Field, IndexRecordOption, STORED, STRING, Schema, TextFieldIndexing, TextOptions, Value,
};
use tantivy::snippet::SnippetGenerator;
use tantivy::tokenizer::TokenStream;
use tantivy::{Index, IndexReader, IndexWriter, TantivyDocument, Term, doc};

use crate::extract::{extract_file, summary_of, supported};
use crate::message::{EXPOSURE_FULL, ResponseItem};
use crate::tokenizer::{TOKENIZER_NAME, analyzer};

const WRITER_BUDGET: usize = 50_000_000;
const MAX_QUERY_TERMS: usize = 10;
const MAX_COMBOS: usize = 256;
const DEFAULT_MIN_COVERAGE: f64 = 0.4;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Collection {
    pub name: String,
    pub path: String,
    pub shared: bool,
    pub exposure: String,
}

#[derive(Debug, Default, Serialize, Deserialize)]
struct CollectionsFile {
    #[serde(default, rename = "collection")]
    collections: Vec<Collection>,
}

pub fn load_collections(manifest: &Path) -> Result<Vec<Collection>> {
    if !manifest.exists() {
        return Ok(Vec::new());
    }
    let raw = fs::read_to_string(manifest).context("reading collections manifest")?;
    let parsed: CollectionsFile = toml::from_str(&raw).context("parsing collections manifest")?;
    Ok(parsed.collections)
}

pub fn save_collections(manifest: &Path, collections: &[Collection]) -> Result<()> {
    if let Some(parent) = manifest.parent() {
        fs::create_dir_all(parent)?;
    }
    let file = CollectionsFile {
        collections: collections.to_vec(),
    };
    fs::write(manifest, toml::to_string_pretty(&file)?)?;
    Ok(())
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchHit {
    pub url: String,
    pub title: String,
    pub summary: String,
    pub published: String,
    pub exposure: String,
    pub collection: String,
    pub shared: bool,
    pub content: Option<String>,
}

pub fn response_items(hits: &[SearchHit]) -> Vec<ResponseItem> {
    hits.iter()
        .map(|hit| ResponseItem {
            url: hit.url.clone(),
            title: hit.title.clone(),
            summary: hit.summary.clone(),
            published: hit.published.clone(),
            exposure: hit.exposure.clone(),
            content: if hit.exposure == EXPOSURE_FULL {
                hit.content.clone()
            } else {
                None
            },
        })
        .collect()
}

struct Fields {
    url: Field,
    title: Field,
    body: Field,
    summary: Field,
    published: Field,
    exposure: Field,
    collection: Field,
    shared: Field,
}

pub struct LocalIndex {
    index: Index,
    reader: IndexReader,
    writer: Mutex<IndexWriter>,
    fields: Fields,
}

impl LocalIndex {
    pub fn open(dir: &Path) -> Result<Self> {
        fs::create_dir_all(dir)?;
        let schema = build_schema();
        let index = Index::open_or_create(MmapDirectory::open(dir)?, schema)
            .context("opening tantivy index")?;
        index.tokenizers().register(TOKENIZER_NAME, analyzer());
        let schema = index.schema();
        let fields = Fields {
            url: schema.get_field("url")?,
            title: schema.get_field("title")?,
            body: schema.get_field("body")?,
            summary: schema.get_field("summary")?,
            published: schema.get_field("published")?,
            exposure: schema.get_field("exposure")?,
            collection: schema.get_field("collection")?,
            shared: schema.get_field("shared")?,
        };
        let reader = index.reader()?;
        let writer = index.writer::<TantivyDocument>(WRITER_BUDGET)?;
        Ok(Self {
            index,
            reader,
            writer: Mutex::new(writer),
            fields,
        })
    }

    pub fn add_collection(&self, collection: &Collection) -> Result<usize> {
        let root = PathBuf::from(&collection.path);
        {
            let writer = self.writer.lock().expect("writer lock");
            writer.delete_term(Term::from_field_text(
                self.fields.collection,
                &collection.name,
            ));
        }
        let mut files = Vec::new();
        walk(&root, &mut files);
        let mut added = 0;
        for file in files {
            if self.add_file(collection, &file)? {
                added += 1;
            }
        }
        self.commit()?;
        Ok(added)
    }

    pub fn add_file(&self, collection: &Collection, path: &Path) -> Result<bool> {
        let Some(extracted) = extract_file(path) else {
            return Ok(false);
        };
        let canonical = path.canonicalize().unwrap_or_else(|_| path.to_path_buf());
        let url = format!("file://{}", canonical.display());
        let summary = summary_of(&extracted.body);
        let shared = if collection.shared { "true" } else { "false" };
        let writer = self.writer.lock().expect("writer lock");
        writer.delete_term(Term::from_field_text(self.fields.url, &url));
        writer.add_document(doc!(
            self.fields.url => url.as_str(),
            self.fields.title => extracted.title.as_str(),
            self.fields.body => extracted.body.as_str(),
            self.fields.summary => summary.as_str(),
            self.fields.published => extracted.published.as_str(),
            self.fields.exposure => collection.exposure.as_str(),
            self.fields.collection => collection.name.as_str(),
            self.fields.shared => shared,
        ))?;
        Ok(true)
    }

    pub fn commit(&self) -> Result<()> {
        self.writer.lock().expect("writer lock").commit()?;
        self.reader.reload()?;
        Ok(())
    }

    pub fn search(
        &self,
        text: &str,
        limit: usize,
        only_shared: bool,
    ) -> Result<(Vec<SearchHit>, u64)> {
        self.search_with_coverage(text, limit, only_shared, DEFAULT_MIN_COVERAGE)
    }

    pub fn search_with_coverage(
        &self,
        text: &str,
        limit: usize,
        only_shared: bool,
        min_coverage: f64,
    ) -> Result<(Vec<SearchHit>, u64)> {
        let limit = limit.max(1);
        self.reader.reload()?;
        let searcher = self.reader.searcher();
        let terms = self.query_terms(text);
        let user_query: Box<dyn Query> = if text.trim().is_empty() {
            Box::new(AllQuery)
        } else if terms.is_empty() {
            Box::new(EmptyQuery)
        } else {
            let gate = self.coverage_query(&terms, min_coverage);
            let mut parts: Vec<(Occur, Box<dyn Query>)> = vec![(Occur::Must, gate)];
            if terms.len() >= 2 {
                let phrase_terms: Vec<(usize, Term)> = terms
                    .iter()
                    .enumerate()
                    .map(|(index, term)| (index, Term::from_field_text(self.fields.body, term)))
                    .collect();
                parts.push((
                    Occur::Should,
                    Box::new(BoostQuery::new(
                        Box::new(PhraseQuery::new_with_offset_and_slop(phrase_terms, 1)),
                        3.0,
                    )),
                ));
            }
            Box::new(BooleanQuery::new(parts))
        };
        let query: Box<dyn Query> = if only_shared {
            let shared_term = TermQuery::new(
                Term::from_field_text(self.fields.shared, "true"),
                IndexRecordOption::Basic,
            );
            Box::new(BooleanQuery::new(vec![
                (Occur::Must, user_query),
                (Occur::Must, Box::new(shared_term)),
            ]))
        } else {
            user_query
        };
        let total = searcher.search(&*query, &Count)? as u64;
        let top = searcher.search(&*query, &TopDocs::with_limit(limit).order_by_score())?;
        let snippet_generator = SnippetGenerator::create(&searcher, &*query, self.fields.body).ok();
        let mut hits = Vec::with_capacity(top.len());
        for (_score, address) in top {
            let document: TantivyDocument = searcher.doc(address)?;
            let snippet = snippet_generator
                .as_ref()
                .map(|generator| generator.snippet_from_doc(&document).fragment().to_string())
                .filter(|fragment| !fragment.trim().is_empty());
            hits.push(SearchHit {
                url: text_value(&document, self.fields.url),
                title: text_value(&document, self.fields.title),
                summary: snippet.unwrap_or_else(|| text_value(&document, self.fields.summary)),
                published: text_value(&document, self.fields.published),
                exposure: text_value(&document, self.fields.exposure),
                collection: text_value(&document, self.fields.collection),
                shared: text_value(&document, self.fields.shared) == "true",
                content: Some(text_value(&document, self.fields.body)),
            });
        }
        Ok((hits, total))
    }

    fn query_terms(&self, text: &str) -> Vec<String> {
        let mut terms = Vec::new();
        let mut seen = HashSet::new();
        if let Some(mut tokenizer) = self.index.tokenizers().get(TOKENIZER_NAME) {
            let mut stream = tokenizer.token_stream(text);
            while let Some(token) = stream.next() {
                let term = token.text.to_string();
                if seen.insert(term.clone()) {
                    terms.push(term);
                }
                if terms.len() >= MAX_QUERY_TERMS {
                    break;
                }
            }
        }
        terms
    }

    fn coverage_query(&self, terms: &[String], min_coverage: f64) -> Box<dyn Query> {
        let n = terms.len();
        let mut required = if n <= 2 {
            n
        } else {
            ((n as f64) * min_coverage).ceil() as usize
        };
        required = required.clamp(1, n);
        let combos = combinations(n, required, MAX_COMBOS);
        let mut shoulds: Vec<(Occur, Box<dyn Query>)> = Vec::with_capacity(combos.len());
        for combo in combos {
            let musts: Vec<(Occur, Box<dyn Query>)> = combo
                .iter()
                .map(|index| {
                    let term = terms[*index].as_str();
                    let term_query: Box<dyn Query> = Box::new(BooleanQuery::new(vec![
                        (
                            Occur::Should,
                            Box::new(BoostQuery::new(
                                Box::new(TermQuery::new(
                                    Term::from_field_text(self.fields.title, term),
                                    IndexRecordOption::WithFreqs,
                                )),
                                2.0,
                            )),
                        ),
                        (
                            Occur::Should,
                            Box::new(TermQuery::new(
                                Term::from_field_text(self.fields.body, term),
                                IndexRecordOption::WithFreqs,
                            )),
                        ),
                    ]));
                    (Occur::Must, term_query)
                })
                .collect();
            shoulds.push((Occur::Should, Box::new(BooleanQuery::new(musts))));
        }
        Box::new(BooleanQuery::new(shoulds))
    }

    pub fn doc_count(&self) -> u64 {
        self.reader.searcher().num_docs()
    }
}

fn binomial(n: usize, k: usize) -> usize {
    if k > n {
        return 0;
    }
    let k = k.min(n - k);
    let mut result = 1usize;
    for index in 0..k {
        result = result * (n - index) / (index + 1);
    }
    result
}

fn combinations(n: usize, start_required: usize, cap: usize) -> Vec<Vec<usize>> {
    let mut required = start_required;
    while required < n && binomial(n, required) > cap {
        required += 1;
    }
    fn walk(
        start: usize,
        remaining: usize,
        n: usize,
        current: &mut Vec<usize>,
        out: &mut Vec<Vec<usize>>,
    ) {
        if remaining == 0 {
            out.push(current.clone());
            return;
        }
        for index in start..=n - remaining {
            current.push(index);
            walk(index + 1, remaining - 1, n, current, out);
            current.pop();
        }
    }
    let mut out = Vec::new();
    walk(0, required, n, &mut Vec::new(), &mut out);
    out
}

fn text_value(document: &TantivyDocument, field: Field) -> String {
    document
        .get_first(field)
        .and_then(|value| value.as_str())
        .unwrap_or("")
        .to_string()
}

fn build_schema() -> Schema {
    let mut builder = Schema::builder();
    let frx_options = TextOptions::default()
        .set_indexing_options(
            TextFieldIndexing::default()
                .set_tokenizer(TOKENIZER_NAME)
                .set_index_option(IndexRecordOption::WithFreqsAndPositions),
        )
        .set_stored();
    builder.add_text_field("url", STRING | STORED);
    builder.add_text_field("title", frx_options.clone());
    builder.add_text_field("body", frx_options);
    builder.add_text_field("summary", STORED);
    builder.add_text_field("published", STORED);
    builder.add_text_field("exposure", STRING | STORED);
    builder.add_text_field("collection", STRING | STORED);
    builder.add_text_field("shared", STRING | STORED);
    builder.build()
}

fn walk(root: &Path, out: &mut Vec<PathBuf>) {
    let Ok(entries) = fs::read_dir(root) else {
        return;
    };
    for entry in entries.flatten() {
        let Ok(file_type) = entry.file_type() else {
            continue;
        };
        if file_type.is_symlink() {
            continue;
        }
        let name = entry.file_name();
        let name = name.to_string_lossy();
        if file_type.is_dir() {
            if name.starts_with('.') {
                continue;
            }
            walk(&entry.path(), out);
        } else if file_type.is_file() && supported(&entry.path()) {
            out.push(entry.path());
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::message::EXPOSURE_METADATA;

    fn collection(name: &str, root: &Path, shared: bool, exposure: &str) -> Collection {
        Collection {
            name: name.to_string(),
            path: root.display().to_string(),
            shared,
            exposure: exposure.to_string(),
        }
    }

    #[test]
    fn shared_filter_and_exposure_are_enforced() {
        let temp = tempfile::tempdir().unwrap();
        let shared_dir = temp.path().join("shared");
        let private_dir = temp.path().join("private");
        fs::create_dir_all(&shared_dir).unwrap();
        fs::create_dir_all(&private_dir).unwrap();
        fs::write(shared_dir.join("a.txt"), "rust ownership rules").unwrap();
        fs::write(private_dir.join("b.txt"), "rust secret notes").unwrap();

        let index = LocalIndex::open(&temp.path().join("index")).unwrap();
        index
            .add_collection(&collection("shared", &shared_dir, true, EXPOSURE_FULL))
            .unwrap();
        index
            .add_collection(&collection(
                "private",
                &private_dir,
                false,
                EXPOSURE_METADATA,
            ))
            .unwrap();

        let (hits, total) = index.search("rust", 10, false).unwrap();
        assert_eq!(total, 2);
        assert_eq!(hits.len(), 2);

        let (shared_hits, shared_total) = index.search("rust", 10, true).unwrap();
        assert_eq!(shared_total, 1);
        assert_eq!(shared_hits.len(), 1);
        assert!(shared_hits[0].url.contains("shared"));

        let items = response_items(&shared_hits);
        assert!(items[0].content.is_some());
    }

    #[test]
    fn coverage_gate_rejects_weak_single_term_overlap() {
        let temp = tempfile::tempdir().unwrap();
        let dir = temp.path().join("corpus");
        fs::create_dir_all(&dir).unwrap();
        fs::write(dir.join("errors.txt"), "error codes and diagnostics").unwrap();
        let index = LocalIndex::open(&temp.path().join("index")).unwrap();
        index
            .add_collection(&collection("test", &dir, true, EXPOSURE_FULL))
            .unwrap();

        let (hits, total) = index.search("launch codes", 10, true).unwrap();
        assert_eq!(total, 0, "two-term queries require both terms");
        assert!(hits.is_empty());
    }

    #[test]
    fn coverage_gate_allows_partial_long_queries() {
        let temp = tempfile::tempdir().unwrap();
        let dir = temp.path().join("corpus");
        fs::create_dir_all(&dir).unwrap();
        fs::write(
            dir.join("manual.txt"),
            "LG DLEX5555 dryer circuit diagram and wiring",
        )
        .unwrap();
        let index = LocalIndex::open(&temp.path().join("index")).unwrap();
        index
            .add_collection(&collection("test", &dir, true, EXPOSURE_FULL))
            .unwrap();

        let (hits, total) = index
            .search("LG dryer 5555 circuit diagram", 10, true)
            .unwrap();
        assert_eq!(total, 1, "four of five terms is enough");
        assert_eq!(hits.len(), 1);
    }

    #[test]
    fn reindex_removes_deleted_files() {
        let temp = tempfile::tempdir().unwrap();
        let dir = temp.path().join("corpus");
        fs::create_dir_all(&dir).unwrap();
        fs::write(dir.join("a.txt"), "rust one").unwrap();
        fs::write(dir.join("b.txt"), "rust two").unwrap();
        let index = LocalIndex::open(&temp.path().join("index")).unwrap();
        index
            .add_collection(&collection("test", &dir, true, EXPOSURE_FULL))
            .unwrap();
        assert_eq!(index.search("rust", 10, false).unwrap().1, 2);

        fs::remove_file(dir.join("b.txt")).unwrap();
        index
            .add_collection(&collection("test", &dir, true, EXPOSURE_FULL))
            .unwrap();
        assert_eq!(index.search("rust", 10, false).unwrap().1, 1);
    }

    #[test]
    fn metadata_exposure_withholds_content() {
        let temp = tempfile::tempdir().unwrap();
        let dir = temp.path().join("corpus");
        fs::create_dir_all(&dir).unwrap();
        fs::write(dir.join("a.txt"), "metadata only").unwrap();
        let index = LocalIndex::open(&temp.path().join("index")).unwrap();
        index
            .add_collection(&collection("test", &dir, true, EXPOSURE_METADATA))
            .unwrap();
        let (hits, _) = index.search("metadata", 10, true).unwrap();
        let items = response_items(&hits);
        assert!(items[0].content.is_none());
        assert_eq!(items[0].exposure, EXPOSURE_METADATA);
    }
}