macltok

use std::{collections::HashMap, path::PathBuf, sync::Arc};

use sha2::{Digest, Sha256};

use tantivy::{

    DocAddress, DocId, DocSet, Index, IndexSettings, ReloadPolicy, Score, Searcher, SegmentReader,

    TantivyDocument, Term,

    collector::{Collector, SegmentCollector},

    directory::{Directory, OwnedBytes, RamDirectory},

    fastfield::{AliveBitSet, write_alive_bitset},

    indexer::{SingleSegmentIndexWriter, merge_filtered_segments},

    query::{BooleanQuery, Occur, Query, QueryParser, TermSetQuery},

    schema::{

        FAST, Field, IndexRecordOption, STORED, STRING, Schema, TextFieldIndexing, TextOptions,

        Value,

    },

    tokenizer::{

        AsciiFoldingFilter, LowerCaser, RemoveLongFilter, SimpleTokenizer, TextAnalyzer,

        TokenStream,

    },

};

use tantivy_common::BitSet;

use super::MessageIndex;

use crate::{Error, ErrorKind};

const TOKENIZER: &str = "macltok";

const WRITER_BUDGET: usize = 4_000_000;

/// segments of a size class merge once there are this many of them

const MERGE_FACTOR: usize = 6;

#[derive(Clone, Debug, Eq, PartialEq)]

pub(crate) struct Doc {

    pub(crate) key: String,

    pub(crate) room: String,

    pub(crate) sender: String,

    /// the original's, or the edit's while the original hasn't turned up

    pub(crate) ts: u64,

    /// the event whose text is indexed, the original or its latest edit

    pub(crate) content: String,

    pub(crate) content_ts: u64,

    /// only kept for provisional docs, read back empty otherwise

    pub(crate) text: String,

    /// an edit seen before its original

    pub(crate) provisional: bool,

}

#[derive(Clone, Debug, Eq, PartialEq)]

pub(crate) struct Hit {

    pub(crate) key: String,

    pub(crate) room: String,

    pub(crate) sender: String,

    pub(crate) content: String,

    pub(crate) ts: u64,

    pub(crate) tie: u64,

}

struct SegmentRecord {

    id: u64,

    del_gen: u64,

}

#[derive(Default)]

pub(crate) struct Changes {

    pub(crate) puts: Vec<(String, Vec<u8>)>,

    pub(crate) deletes: Vec<String>,

    files: Vec<u64>,

    dels: Vec<(u64, u64)>,

    removed: Vec<u64>,

}

#[derive(Clone, Copy)]

struct Fields {

    key: Field,

    room: Field,

    sender: Field,

    content: Field,

    text: Field,

    ts: Field,

    tie: Field,

    content_ts: Field,

    pending: Field,

}

struct Segment {

    id: u64,

    raw: RamDirectory,

    index: Index,

    searcher: Searcher,

    alive: Arc<BitSet>,

    max_doc: u32,

    del_gen: u64,

    persisted: bool,

    deletes_dirty: bool,

}

impl Segment {

    fn live(&self) -> usize {

        self.alive.len()

    }

}

pub(crate) struct KeywordIndex {

    schema: Schema,

    fields: Fields,

    settings: IndexSettings,

    parser: QueryParser,

    segments: Vec<Segment>,

    buffer: HashMap<String, Doc>,

    next_id: u64,

    persistent: bool,

    /// merged away since the last commit, with the delete generation they were persisted at

    removed: Vec<(u64, bool, u64)>,

}

pub(crate) fn analyzer() -> TextAnalyzer {

    TextAnalyzer::builder(SimpleTokenizer::default())

        .filter(RemoveLongFilter::limit(40))

        .filter(LowerCaser)

        .filter(AsciiFoldingFilter)

        .build()

}

/// the byte ranges in `text` whose analysed token is one of `terms`

pub(crate) fn match_ranges(text: &str, terms: &[String]) -> Vec<(usize, usize)> {

    if terms.is_empty() {

        return Vec::new();

    }

    let mut analyzer = analyzer();

    let mut stream = analyzer.token_stream(text);

    let mut ranges = Vec::new();

    while stream.advance() {

        let token = stream.token();

        if terms.iter().any(|term| *term == token.text) {

            ranges.push((token.offset_from, token.offset_to));

        }

    }

    ranges

}

fn schema() -> (Schema, Fields) {

    let mut builder = Schema::builder();

    let text_options = TextOptions::default().set_indexing_options(

        TextFieldIndexing::default()

            .set_tokenizer(TOKENIZER)

            .set_index_option(IndexRecordOption::WithFreqsAndPositions),

    );

    let fields = Fields {

        key: builder.add_text_field("key", STRING | STORED),

        room: builder.add_text_field("room", STRING | STORED),

        sender: builder.add_text_field("sender", STRING | STORED),

        content: builder.add_text_field("content", STRING | STORED),

        text: builder.add_text_field("text", text_options),

        ts: builder.add_u64_field("ts", FAST | STORED),

        tie: builder.add_u64_field("tie", FAST),

        content_ts: builder.add_u64_field("content_ts", STORED),

        pending: builder.add_text_field("pending", STORED),

    };

    (builder.build(), fields)

}

fn tie(key: &str) -> u64 {

    let digest = Sha256::digest(key.as_bytes());

    u64::from_le_bytes(digest[..8].try_into().expect("a sha256 digest is 32 bytes"))

}

fn settings() -> IndexSettings {

    // the dedicated compressor is a thread, and wasm32-unknown-unknown has none

    IndexSettings {

        docstore_compress_dedicated_thread: false,

        ..IndexSettings::default()

    }

}

fn tantivy_error(

    context: &'static str,

    error: impl std::error::Error + Send + Sync + 'static,

) -> Error {

    Error::with_source(ErrorKind::Store, context, error)

}

fn open_segment(id: u64, raw: RamDirectory, max_doc_hint: Option<u32>) -> Result<Segment, Error> {

    let index = Index::open(raw.clone())

        .map_err(|error| tantivy_error("opening a search segment", error))?;

    index.tokenizers().register(TOKENIZER, analyzer());

    let reader = index

        .reader_builder()

        .reload_policy(ReloadPolicy::Manual)

        .try_into()

        .map_err(|error| tantivy_error("reading a search segment", error))?;

    let searcher = reader.searcher();

    let max_doc = searcher

        .segment_readers()

        .iter()

        .map(SegmentReader::max_doc)

        .sum();

    if searcher.segment_readers().len() > 1 || max_doc_hint.is_some_and(|hint| hint != max_doc) {

        return Err(Error::new(

            ErrorKind::Store,

            "a search segment has an unexpected shape",

        ));

    }

    Ok(Segment {

        id,

        raw,

        index,

        searcher,

        alive: Arc::new(BitSet::with_max_value_and_full(max_doc)),

        max_doc,

        del_gen: 0,

        persisted: false,

        deletes_dirty: false,

    })

}

fn pack(files: &[(String, Vec<u8>)]) -> Vec<u8> {

    let mut packed = Vec::with_capacity(

        files

            .iter()

            .map(|(name, bytes)| name.len() + bytes.len() + 12)

            .sum(),

    );

    packed.extend_from_slice(&(files.len() as u32).to_le_bytes());

    for (name, bytes) in files {

        packed.extend_from_slice(&(name.len() as u32).to_le_bytes());

        packed.extend_from_slice(name.as_bytes());

        packed.extend_from_slice(&(bytes.len() as u64).to_le_bytes());

        packed.extend_from_slice(bytes);

    }

    packed

}

fn unpack(mut bytes: &[u8]) -> Option<Vec<(String, Vec<u8>)>> {

    fn take<'a>(bytes: &mut &'a [u8], len: usize) -> Option<&'a [u8]> {

        if bytes.len() < len {

            return None;

        }

        let (head, tail) = bytes.split_at(len);

        *bytes = tail;

        Some(head)

    }

    let count = u32::from_le_bytes(take(&mut bytes, 4)?.try_into().ok()?);

    let mut files = Vec::new();

    for _ in 0..count {

        let name_len = u32::from_le_bytes(take(&mut bytes, 4)?.try_into().ok()?) as usize;

        let name = std::str::from_utf8(take(&mut bytes, name_len)?)

            .ok()?

            .to_owned();

        let len =

            usize::try_from(u64::from_le_bytes(take(&mut bytes, 8)?.try_into().ok()?)).ok()?;

        files.push((name, take(&mut bytes, len)?.to_vec()));

    }

    bytes.is_empty().then_some(files)

}

pub(crate) fn segment_key(id: u64) -> String {

    format!("seg/{id}")

}

pub(crate) fn deletes_key(id: u64, generation: u64) -> String {

    format!("del/{id}/{generation}")

}

impl KeywordIndex {

    pub(crate) fn new(persistent: bool) -> Self {

        let (schema, fields) = schema();

        let parser_index = Index::create_in_ram(schema.clone());

        parser_index.tokenizers().register(TOKENIZER, analyzer());

        let mut parser = QueryParser::for_index(&parser_index, vec![fields.text]);

        parser.set_conjunction_by_default();

        Self {

            schema,

            fields,

            settings: settings(),

            parser,

            segments: Vec::new(),

            buffer: HashMap::new(),

            next_id: 1,

            persistent,

            removed: Vec::new(),

        }

    }

    /// `blob` hands back what an earlier commit put, any of it missing or unreadable fails the load

    pub(crate) fn load(

        manifest: &serde_json::Value,

        mut blob: impl FnMut(&str) -> Option<Vec<u8>>,

    ) -> Result<Self, Error> {

        let invalid = || Error::new(ErrorKind::Store, "the search manifest is unreadable");

        let records = manifest["segments"]

            .as_array()

            .ok_or_else(invalid)?

            .iter()

            .map(|record| {

                Some(SegmentRecord {

                    id: record.get(0)?.as_u64()?,

                    del_gen: record.get(1)?.as_u64()?,

                })

            })

            .collect::<Option<Vec<_>>>()

            .ok_or_else(invalid)?;

        let next_id = manifest["next"].as_u64().ok_or_else(invalid)?;

        let mut index = Self::new(true);

        let missing = || {

            Error::new(

                ErrorKind::Store,

                "a search segment is missing or unreadable",

            )

        };

        for record in &records {

            let files = blob(&segment_key(record.id))

                .and_then(|bytes| unpack(&bytes))

                .ok_or_else(missing)?;

            let raw = RamDirectory::create();

            for (name, bytes) in files {

                raw.atomic_write(&PathBuf::from(name), &bytes)

                    .map_err(|error| tantivy_error("loading a search segment", error))?;

            }

            let mut segment = open_segment(record.id, raw, None)?;

            if segment.index.schema() != index.schema {

                return Err(Error::new(

                    ErrorKind::Store,

                    "a search segment has another schema",

                ));

            }

            if record.del_gen > 0 {

                let bytes = blob(&deletes_key(record.id, record.del_gen)).ok_or_else(missing)?;

                let alive = AliveBitSet::open(OwnedBytes::new(bytes));

                let mut bits = BitSet::with_max_value(segment.max_doc);

                for doc in alive.iter_alive().filter(|doc| *doc < segment.max_doc) {

                    bits.insert(doc);

                }

                segment.alive = Arc::new(bits);

            }

            segment.del_gen = record.del_gen;

            segment.persisted = true;

            index.segments.push(segment);

        }

        index.next_id = next_id.max(

            records

                .iter()

                .map(|record| record.id + 1)

                .max()

                .unwrap_or(1),

        );

        Ok(index)

    }

    fn records(&self) -> Vec<SegmentRecord> {

        self.segments

            .iter()

            .filter(|segment| segment.persisted)

            .map(|segment| SegmentRecord {

                id: segment.id,

                del_gen: segment.del_gen,

            })

            .collect()

    }

    fn document(&self, doc: &Doc) -> TantivyDocument {

        let fields = self.fields;

        let mut document = TantivyDocument::default();

        document.add_text(fields.key, &doc.key);

        document.add_text(fields.room, &doc.room);

        document.add_text(fields.sender, &doc.sender);

        document.add_text(fields.content, &doc.content);

        document.add_text(fields.text, &doc.text);

        document.add_u64(fields.ts, doc.ts);

        document.add_u64(fields.tie, tie(&doc.key));

        document.add_u64(fields.content_ts, doc.content_ts);

        if doc.provisional {

            document.add_text(fields.pending, &doc.text);

        }

        document

    }

    fn stored(&self, document: &TantivyDocument) -> Option<Doc> {

        let fields = self.fields;

        let text = |field| {

            document

                .get_first(field)

                .and_then(|value| value.as_str())

                .map(str::to_owned)

        };

        let number = |field| document.get_first(field).and_then(|value| value.as_u64());

        let pending = text(fields.pending);

        Some(Doc {

            key: text(fields.key)?,

            room: text(fields.room)?,

            sender: text(fields.sender)?,

            ts: number(fields.ts)?,

            content: text(fields.content)?,

            content_ts: number(fields.content_ts)?,

            provisional: pending.is_some(),

            text: pending.unwrap_or_default(),

        })

    }

    /// the live docs in `segment` whose `field` holds `value`

    fn find(&self, segment: &Segment, field: Field, value: &str) -> Vec<DocId> {

        let term = Term::from_field_text(field, value);

        let mut found = Vec::new();

        for reader in segment.searcher.segment_readers() {

            let Ok(inverted) = reader.inverted_index(field) else {

                continue;

            };

            let Ok(Some(mut postings)) = inverted.read_postings(&term, IndexRecordOption::Basic)

            else {

                continue;

            };

            let mut doc = postings.doc();

            while doc != tantivy::TERMINATED {

                if segment.alive.contains(doc) {

                    found.push(doc);

                }

                doc = postings.advance();

            }

        }

        found

    }

    fn lookup(&self, field: Field, value: &str) -> Option<Doc> {

        self.segments.iter().rev().find_map(|segment| {

            let doc = *self.find(segment, field, value).first()?;

            let document: TantivyDocument = segment.searcher.doc(DocAddress::new(0, doc)).ok()?;

            self.stored(&document)

        })

    }

    fn flush_buffer(&mut self) -> Result<(), Error> {

        if self.buffer.is_empty() {

            return Ok(());

        }

        let raw = RamDirectory::create();

        let index = Index::create(raw.clone(), self.schema.clone(), self.settings.clone())

            .map_err(|error| tantivy_error("creating a search segment", error))?;

        index.tokenizers().register(TOKENIZER, analyzer());

        let mut writer = SingleSegmentIndexWriter::new(index, WRITER_BUDGET)

            .map_err(|error| tantivy_error("writing a search segment", error))?;

        let mut docs: Vec<_> = self.buffer.values().collect();

        docs.sort_by_key(|doc| doc.ts);

        for doc in docs {

            writer

                .add_document(self.document(doc))

                .map_err(|error| tantivy_error("writing a search segment", error))?;

        }

        writer

            .finalize()

            .map_err(|error| tantivy_error("finishing a search segment", error))?;

        let segment = open_segment(self.next_id, raw, Some(self.buffer.len() as u32))?;

        self.next_id += 1;

        self.segments.push(segment);

        self.buffer.clear();

        Ok(())

    }

    fn class(live: usize) -> u32 {

        live.max(1).ilog(MERGE_FACTOR)

    }

    /// merges a size class once it fills up, and rewrites a segment that's mostly deleted

    fn compact(&mut self) -> Result<(), Error> {

        self.segments.retain(|segment| {