macltok

mod feed;

#[cfg_attr(not(any(feature = "sqlite", feature = "indexeddb")), allow(dead_code))]

mod index;

#[cfg(any(feature = "sqlite", feature = "indexeddb"))]

pub(crate) mod store;

use std::{

    collections::{BTreeMap, HashMap, HashSet, VecDeque},

    sync::Arc,

    time::Duration,

};

use futures_util::{

    StreamExt,

    future::join_all,

    stream::{AbortHandle, FuturesUnordered, SelectAll, abortable},

};

use matrix_sdk::{

    Client, Room, RoomState, SendOutsideWasm,

    deserialized_responses::TimelineEvent as SdkTimelineEvent,

    event_cache::{RoomEventCacheSubscriber, RoomEventCacheUpdate},

    room::MessagesOptions,

};

use matrix_sdk_base::RoomInfoNotableUpdateReasons;

use matrix_sdk_ui::eyeball_im::VectorDiff;

use ruma::{

    OwnedRoomId, OwnedUserId, UInt,

    events::{

        ignored_user_list::IgnoredUserListEventContent,

        room::encrypted::OriginalSyncRoomEncryptedEvent,

    },

    serde::Raw,

};

use serde_json::Value;

use tokio::sync::{Mutex, broadcast, mpsc, oneshot};

use self::{

    feed::{Effect, Observed, Redactions},

    index::{Changes, Doc, Hit, KeywordIndex},

};

use crate::{

    CommandOutput, EngineConfig, Error, ErrorKind, RoomSearchCoverage, SearchBackfill,

    SearchCommand, SearchCoverage, SearchHit, SearchMatch, SearchPage, SearchScope, SearchSession,

    engine::{MediaService, SearchStore, context, executor::EngineBoxFuture},

    normalise::{MediaRegistration, MediaScope},

};

const MAX_LIMIT: u16 = 50;

/// a consumer that never closes its searches still holds at most this many

const MAX_SEARCHES: usize = 8;

const COMMIT_AFTER_MS: u64 = 5_000;

const COMMIT_BUFFERED: usize = 2_000;

#[cfg(any(feature = "sqlite", feature = "indexeddb"))]

const COMMIT_RETRY_MS: u64 = 60_000;

const IDLE_SUBSCRIPTION_MS: u64 = 5 * 60_000;

const BACKFILL_RETRY_MS: u64 = 60_000;

const BACKFILL_RETRY_MAX_MS: u64 = 15 * 60_000;

#[cfg(any(feature = "sqlite", feature = "indexeddb"))]

const SCHEMA: u64 = 1;

#[cfg(not(target_family = "wasm"))]

type BoxStream<T> = futures_util::stream::BoxStream<'static, T>;

#[cfg(target_family = "wasm")]

type BoxStream<T> = futures_util::stream::LocalBoxStream<'static, T>;

/// a local index of message text. keyword search over tantivy is the only one today, an

/// embedding index can stand beside it later without the public api noticing

#[cfg_attr(not(any(feature = "sqlite", feature = "indexeddb")), allow(dead_code))]

pub(crate) trait MessageIndex: SendOutsideWasm {

    fn get(&self, key: &str) -> Option<Doc>;

    fn by_content(&self, content: &str) -> Option<Doc>;

    /// replaces whatever `doc.key` held

    fn put(&mut self, doc: Doc);

    fn remove(&mut self, key: &str) -> bool;

    fn buffered(&self) -> usize;

    /// newest first, strictly older than `before`. also hands back the analysed query terms

    fn search(

        &mut self,

        query: &str,

        rooms: Option<&[String]>,

        before: Option<(u64, u64)>,

        limit: usize,

    ) -> Result<(Vec<Hit>, Vec<String>), Error>;

    fn maintain(&mut self) -> Result<(), Error>;

    fn changes(&mut self) -> Result<Changes, Error>;

    fn committed(&mut self, changes: &Changes);

    fn manifest(&self) -> Value;

}

#[derive(Clone, Debug, Default)]

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]

struct RoomProgress {

    oldest: Option<u64>,

    newest: Option<u64>,

    /// where the next backfill page starts, `None` before the first one

    tail: Option<String>,

    started: bool,

    start_reached: bool,

}

#[derive(Clone, Debug)]

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]

struct Pending {

    room: String,

    session_id: String,

    ts: u64,

    raw: String,

}

#[derive(Clone, Copy)]

struct Budget {

    pages_per_room: u16,

    page_size: u16,

    interval_ms: u64,

    max_pending: usize,

}

impl Budget {

    fn from_config(config: &EngineConfig) -> Self {

        Self {

            pages_per_room: config.search_backfill_pages_per_room,

            page_size: config.search_backfill_page_size.max(1),

            interval_ms: config.search_backfill_interval_ms,

            max_pending: config.max_search_undecrypted,

        }

    }

}

struct RawPage {

    hits: Vec<Hit>,

    terms: Vec<String>,

    more: bool,

    coverage: SearchCoverage,

}

enum Request {

    Search {

        query: String,

        rooms: Option<Vec<String>>,

        before: Option<(u64, u64)>,

        limit: usize,

        coverage: Vec<OwnedRoomId>,

        reply: oneshot::Sender<Result<RawPage, Error>>,

    },

    Coverage(Vec<OwnedRoomId>, oneshot::Sender<SearchCoverage>),

    /// a hit that turned out redacted when its preview was read

    Forget(String),

    Stop(oneshot::Sender<()>),

}

#[derive(Clone)]

pub(crate) struct Handle {

    requests: mpsc::Sender<Request>,

    client: Client,

    media: Arc<MediaService>,

    sessions: Arc<Sessions>,

}

enum Work {

    Backfill {

        room: OwnedRoomId,

        head: bool,

        result: Result<matrix_sdk::room::Messages, Error>,

    },

    Decrypted {

        key: String,

        room: OwnedRoomId,

        value: Option<Value>,

    },

    Refetched {

        room: OwnedRoomId,

        value: Option<Value>,

    },

}

struct Failure {

    message: String,

    retry_at_ms: u64,

}

pub(crate) struct Service {

    client: Client,

    index: Box<dyn MessageIndex>,

    store: Option<Result<SearchStore, Error>>,

    rooms: BTreeMap<String, RoomProgress>,

    pending: BTreeMap<String, Pending>,

    pending_dirty: bool,

    redactions: Redactions,

    dirty_since: Option<u64>,

    commit_after: u64,

    budget: Budget,

    requests: mpsc::Receiver<Request>,

    subscriptions: HashMap<OwnedRoomId, (AbortHandle, u64)>,

    updates: SelectAll<BoxStream<(OwnedRoomId, Option<Vec<SdkTimelineEvent>>)>>,

    work: FuturesUnordered<EngineBoxFuture<'static, Work>>,

    pages: HashMap<OwnedRoomId, u16>,

    heads: VecDeque<OwnedRoomId>,

    headed: HashSet<OwnedRoomId>,

    head_tokens: HashMap<OwnedRoomId, String>,

    backfilling: Option<OwnedRoomId>,

    next_backfill: u64,

    retry_delay: u64,

    failures: HashMap<OwnedRoomId, Failure>,

    turn: usize,

}

fn now() -> u64 {

    ruma::MilliSecondsSinceUnixEpoch::now().get().into()

}

fn value_of(event: &SdkTimelineEvent) -> Option<Value> {

    serde_json::from_str(event.raw().json().get()).ok()

}

fn encode_cursor((ts, tie): (u64, u64)) -> String {

    format!("{ts:x}.{tie:x}")

}

fn decode_cursor(cursor: &str) -> Option<(u64, u64)> {

    let (ts, tie) = cursor.split_once('.')?;

    Some((

        u64::from_str_radix(ts, 16).ok()?,

        u64::from_str_radix(tie, 16).ok()?,

    ))

}

impl Handle {

    pub(crate) async fn execute(&self, command: SearchCommand) -> Result<CommandOutput, Error> {

        match command {

            SearchCommand::Search {

                query,

                scope,

                limit,

            } => {

                let (page, media) = self.search(&query, &scope, None, limit).await?;

                let search_id = self.sessions.open(&self.media, query, scope, media).await;

                Ok(CommandOutput::SearchSession(SearchSession { search_id, page }))

            }

            SearchCommand::More {

                search_id,

                cursor,

                limit,

            } => {

                let (query, scope) = self.sessions.get(&search_id).await?;

                let (page, media) = self.search(&query, &scope, Some(cursor.as_str()), limit).await?;

                self.sessions.register(&self.media, &search_id, media).await?;

                Ok(CommandOutput::SearchPage(page))

            }

            SearchCommand::Close { search_id } => {

                self.sessions.close(&self.media, &search_id).await?;

                Ok(CommandOutput::Done)

            }

            SearchCommand::Coverage(scope) => {

                let (reply, response) = oneshot::channel();

                self.requests

                    .send(Request::Coverage(self.scope_rooms(&scope), reply))

                    .await

                    .map_err(|_| Error::interrupted())?;

                Ok(CommandOutput::SearchCoverage(

                    response.await.map_err(|_| Error::interrupted())?,

                ))

            }

        }

    }

    /// the media service is cleared beside this, so nothing is released one by one

    pub(crate) async fn stop(&self) {

        let (reply, response) = oneshot::channel();

        if self.requests.send(Request::Stop(reply)).await.is_ok() {

            let _ = response.await;

        }

        self.sessions.open.lock().await.clear();

    }

    fn scope_rooms(&self, scope: &SearchScope) -> Vec<OwnedRoomId> {

        match scope {

            SearchScope::Room(room_id) => vec![room_id.clone()],

            SearchScope::Rooms(room_ids) => room_ids.clone(),

            SearchScope::All => self

                .client

                .joined_rooms()

                .into_iter()

                .map(|room| room.room_id().to_owned())

                .collect(),

        }

    }

    async fn search(

        &self,

        query: &str,

        scope: &SearchScope,

        cursor: Option<&str>,

        limit: u16,

    ) -> Result<(SearchPage, Vec<MediaRegistration>), Error> {

        if limit == 0 || limit > MAX_LIMIT {

            return Err(Error::new(

                ErrorKind::Invalid,

                "search limit must be between 1 and 50",

            ));

        }

        if query.trim().is_empty() {

            return Err(Error::new(ErrorKind::Invalid, "search query is empty"));

        }

        let before = match cursor {

            Some(cursor) => Some(

                decode_cursor(cursor)

                    .ok_or_else(|| Error::new(ErrorKind::Invalid, "search cursor is invalid"))?,

            ),

            None => None,

        };

        let rooms = self.scope_rooms(scope);

        let filter = Some(rooms.iter().map(ToString::to_string).collect());

        let (reply, response) = oneshot::channel();

        self.requests

            .send(Request::Search {

                query: query.to_owned(),

                rooms: filter,

                before,

                limit: usize::from(limit),

                coverage: rooms,

                reply,

            })

            .await

            .map_err(|_| Error::interrupted())?;

        let raw = response.await.map_err(|_| Error::interrupted())??;

        let next_cursor = raw

            .more

            .then(|| raw.hits.last().map(|hit| encode_cursor((hit.ts, hit.tie))))

            .flatten();

        let ignored = self.ignored_users().await;

        let current_user_id = self.client.user_id().map(ToOwned::to_owned);

        let previews = join_all(

            raw.hits

                .into_iter()

                .filter(|hit| !ignored.iter().any(|user| user.as_str() == hit.sender))

                .map(|hit| self.preview(hit, current_user_id.clone(), &raw.terms)),

        )

        .await;

        let mut hits = Vec::new();

        let mut media = Vec::new();

        for preview in previews {

            match preview {

                Preview::Shown(hit, registrations) => {

                    hits.push(hit);

                    media.extend(registrations);

                }

                Preview::Redacted(key) => {

                    let _ = self.requests.send(Request::Forget(key)).await;

                }

                Preview::Unavailable => {}

            }

        }

        let page = SearchPage {

            hits,

            next_cursor,

            coverage: raw.coverage,

        };

        Ok((page, media))

    }

    async fn ignored_users(&self) -> Vec<OwnedUserId> {

        self.client

            .account()

            .account_data::<IgnoredUserListEventContent>()

            .await

            .ok()

            .flatten()

            .and_then(|raw| raw.deserialize().ok())

            .map(|content| content.ignored_users.into_keys().collect())

            .unwrap_or_default()

    }

    async fn preview(

        &self,

        hit: Hit,

        current_user_id: Option<OwnedUserId>,

        terms: &[String],

    ) -> Preview {

        let (Ok(room_id), Ok(event_id)) = (

            OwnedRoomId::try_from(hit.room.as_str()),

            ruma::OwnedEventId::try_from(hit.key.as_str()),

        ) else {

            return Preview::Unavailable;

        };

        let Some(room) = self.client.get_room(&room_id) else {

            return Preview::Unavailable;

        };

        let Some((original, highlight)) = load(&room, &hit.key).await else {

            return Preview::Unavailable;

        };

        if original.pointer("/unsigned/redacted_because").is_some() {

            return Preview::Redacted(hit.key);

        }

        let mut values = vec![original];

        if hit.content != hit.key

            && let Some((edit, _)) = load(&room, &hit.content).await

        {

            values.push(edit);

        }

        let Some((event, registrations)) =

            context::preview(&room, current_user_id.as_deref(), &values, highlight).await

        else {

            return Preview::Unavailable;

        };

        let text = match &event.content {

            crate::TimelineContent::Message(message) => Some(message.body.as_str()),

            crate::TimelineContent::Sticker { body, .. } => Some(body.as_str()),

            crate::TimelineContent::Poll { question } => Some(question.as_str()),

            _ => None,

        };

        let matches = text

            .map(|text| index::match_ranges(text, terms))

            .unwrap_or_default()

            .into_iter()

            .filter_map(|(start, end)| {

                Some(SearchMatch {

                    start: u32::try_from(start).ok()?,

                    end: u32::try_from(end).ok()?,

                })

            })

            .collect();

        Preview::Shown(

            SearchHit {

                room_id,

                event_id,

                timestamp_ms: event.timestamp_ms,

                event,

                matches,

            },

            registrations,

        )

    }

}

enum Preview {

    Shown(SearchHit, Vec<MediaRegistration>),

    Redacted(String),

    Unavailable,

}

/// open searches, oldest first, each owning the media of the hits it handed out

#[derive(Default)]

struct Sessions {

    open: Mutex<VecDeque<Session>>,

}

struct Session {

    id: String,

    query: String,

    scope: SearchScope,

    media: MediaScope,

}

impl Sessions {

    async fn open(

        &self,

        media: &MediaService,

        query: String,

        scope: SearchScope,

        registrations: Vec<MediaRegistration>,

    ) -> String {

        let media_scope = media.context_scope();

        let id = format!("search-{}", media_scope.nonce());

        let mut open = self.open.lock().await;

        media.register(&media_scope, registrations).await;

        open.push_back(Session {

            id: id.clone(),

            query,

            scope,

            media: media_scope,

        });

        while open.len() > MAX_SEARCHES {

            if let Some(oldest) = open.pop_front() {

                media.release(&oldest.media).await;