macltok

use std::{

    collections::{BTreeMap, BTreeSet},

    sync::{

        Arc,

        atomic::{AtomicBool, Ordering},

    },

    time::Duration,

};

use futures_util::{StreamExt, stream::FuturesUnordered};

use matrix_sdk::{

    Client, Room, RoomState,

    event_cache::RedecryptorReport,

    sliding_sync::{ResponseCommitBatch, ResponseCommitHook},

};

use ruma::{

    OwnedEventId, OwnedRoomId, OwnedUserId, RoomId,

    events::{

        AnySyncTimelineEvent, ignored_user_list::IgnoredUserListEventContent,

        room::encrypted::OriginalSyncRoomEncryptedEvent,

    },

    serde::Raw,

};

use serde_json::Value;

use tokio::sync::{mpsc, oneshot, watch};

use super::state::{

    Candidate, Checkpoint, Data, RepairCoverage, SavedEvent, SyncState, Window, id, invalid, now,

};

use crate::{

    CommandOutput, EngineConfig, Error, ErrorKind, Event, HistoryRepairPolicy, MentionCaptureState,

    MentionEntry, MentionGapReason, MentionKey, MentionReason, MentionsCommand, MentionsPage,

    MentionsUpdate, TimelineContent, TimelineItem,

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

    normalise::{MediaRegistration, normalise_observed},

    transport::EventSender,

};

#[cfg(all(test, feature = "sqlite"))]

#[path = "service_tests.rs"]

mod tests;

const HOOK_TIMEOUT: Duration = Duration::from_secs(10);

const FLUSH_RETRY_MS: u64 = 60_000;

#[derive(Clone)]

pub(crate) struct Handle {

    requests: mpsc::Sender<Request>,

    timed_out: Arc<AtomicBool>,

}

enum Request {

    Batch(ResponseCommitBatch, oneshot::Sender<()>),

    Command(

        MentionsCommand,

        oneshot::Sender<Result<CommandOutput, Error>>,

    ),

    GetRepairPolicy(oneshot::Sender<Result<HistoryRepairPolicy, Error>>),

    SetRepairPolicy(HistoryRepairPolicy, oneshot::Sender<Result<(), Error>>),

    GetRepairCoverage(OwnedRoomId, oneshot::Sender<Result<RepairCoverage, Error>>),

    RepairResolved(OwnedRoomId, oneshot::Sender<Result<(), Error>>),

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

}

pub(crate) struct Service {

    client: Client,

    account: OwnedUserId,

    media: Arc<MediaService>,

    store: Option<MentionsStore>,

    data: Data,

    sync: SyncState,

    dirty: bool,

    flush_after: u64,

    ignored: BTreeSet<OwnedUserId>,

    leaving: BTreeSet<OwnedRoomId>,

    timed_out: Arc<AtomicBool>,

    config: EngineConfig,

    requests: mpsc::Receiver<Request>,

    updates: watch::Sender<MentionsPage>,

    recovery_budget: usize,

    latest_scan: bool,

    recovery_retry: u64,

}

enum Work {

    Capture(

        String,

        u64,

        Result<(crate::EventContext, Vec<MediaRegistration>, Vec<Value>), Error>,

    ),

    Candidate(String, Result<Value, Error>),

    History(

        bool,

        Option<String>,

        Result<ruma::api::client::push::get_notifications::v3::Response, Error>,

    ),

}

impl Handle {

    /// never fails, a mentions failure becomes a coverage gap rather than a stalled sync

    pub(crate) fn hook(&self) -> ResponseCommitHook {

        let handle = self.clone();

        ResponseCommitHook::new(move |batch| {

            let handle = handle.clone();

            async move {

                let committed = async {

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

                    handle

                        .requests

                        .send(Request::Batch(batch, reply))

                        .await

                        .ok()?;

                    response.await.ok()

                };

                match matrix_sdk::timeout::timeout(committed, HOOK_TIMEOUT).await {

                    Ok(Some(())) => {}

                    Ok(None) => tracing::warn!("mentions service stopped, sync carries on"),

                    Err(_) => {

                        handle.timed_out.store(true, Ordering::Release);

                        tracing::warn!("mentions commit timed out, sync carries on");

                    }

                }

                Ok(())

            }

        })

    }

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

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

        self.requests

            .send(Request::Command(command, reply))

            .await

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

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

    }

    pub(crate) async fn repair_policy(&self) -> Result<HistoryRepairPolicy, Error> {

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

        self.requests

            .send(Request::GetRepairPolicy(reply))

            .await

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

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

    }

    pub(crate) async fn set_repair_policy(&self, policy: HistoryRepairPolicy) -> Result<(), Error> {

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

        self.requests

            .send(Request::SetRepairPolicy(policy, reply))

            .await

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

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

    }

    pub(crate) async fn repair_coverage(

        &self,

        room_id: OwnedRoomId,

    ) -> Result<RepairCoverage, Error> {

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

        self.requests

            .send(Request::GetRepairCoverage(room_id, reply))

            .await

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

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

    }

    pub(crate) async fn repair_resolved(&self, room_id: OwnedRoomId) -> Result<(), Error> {

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

        self.requests

            .send(Request::RepairResolved(room_id, reply))

            .await

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

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

    }

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

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

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

            let _ = response.await;

        }

    }

}

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

async fn open_store(

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

) -> (

    Option<MentionsStore>,

    Data,

    SyncState,

    Vec<MentionGapReason>,

) {

    let mut store = match store {

        None => return (None, Data::default(), SyncState::default(), Vec::new()),

        Some(Err(error)) => {

            tracing::error!(%error, "mentions store unavailable, the inbox runs in memory");

            return (

                None,

                Data::default(),

                SyncState::default(),

                vec![MentionGapReason::StoreFailed],

            );

        }

        Some(Ok(store)) => store,

    };

    let loaded = (|| {

        let data: Data = store.read("inbox")?.unwrap_or_default();

        let sync: SyncState = store.read("sync")?.unwrap_or_default();

        if sync.inbox_revision != data.revision {

            return Err(Error::new(

                ErrorKind::Store,

                "mentions checkpoint does not match its inbox",

            ));

        }

        Ok((data, sync))

    })();

    let (data, sync) = match loaded {

        Ok(loaded) => loaded,

        Err(error) => {

            tracing::warn!(%error, "resetting an unreadable mentions inbox");

            if let Err(error) = store.reset().await {

                tracing::error!(%error, "mentions store unavailable, the inbox runs in memory");

                return (

                    None,

                    Data::default(),

                    SyncState::default(),

                    vec![MentionGapReason::StoreFailed],

                );

            }

            (Data::default(), SyncState::default())

        }

    };

    let gaps = if store.was_reset() {

        vec![MentionGapReason::StoreReset]

    } else {

        Vec::new()

    };

    (Some(store), data, sync, gaps)

}

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

async fn open_store(

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

) -> (

    Option<MentionsStore>,

    Data,

    SyncState,

    Vec<MentionGapReason>,

) {

    let gaps = match store {

        None => Vec::new(),

        Some(Err(_)) => vec![MentionGapReason::StoreFailed],

        Some(Ok(store)) => match store {},

    };

    (None, Data::default(), SyncState::default(), gaps)

}

impl Service {

    pub(crate) async fn new(

        client: Client,

        media: Arc<MediaService>,

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

        config: EngineConfig,

    ) -> Result<(Handle, Self, watch::Receiver<MentionsPage>), Error> {

        let account = client

            .user_id()

            .ok_or_else(|| {

                Error::new(

                    ErrorKind::Unauthenticated,

                    "mentions require an authenticated account",

                )

            })?

            .to_owned();

        let (store, mut data, sync, gaps) = open_store(store).await;

        for reason in &gaps {

            data.gap(None, *reason, None);

        }

        let page = data.page(&account, None, 50)?;

        media

            .register(

                page.entries

                    .iter()

                    .filter_map(|entry| data.events.get(&id(&entry.key)))

                    .flat_map(|event| event.media.clone())

                    .collect(),

            )

            .await;

        let (updates, receiver) = watch::channel(page);

        let (sender, requests) = mpsc::channel(config.command_capacity.max(1));

        let timed_out = Arc::new(AtomicBool::new(false));

        let mut service = Self {

            client,

            account,

            media,

            store,

            data,

            sync,

            dirty: !gaps.is_empty(),

            flush_after: 0,

            ignored: BTreeSet::new(),

            leaving: BTreeSet::new(),

            timed_out: timed_out.clone(),

            config,

            requests,

            updates,

            recovery_budget: 20,

            latest_scan: true,

            recovery_retry: 0,

        };

        service.ignored = service.ignored_users().await;

        Ok((

            Handle {

                requests: sender,

                timed_out,

            },

            service,

            receiver,

        ))

    }

    async fn ignored_users(&self) -> BTreeSet<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 publish(&mut self) {

        let Ok(page) = self.data.page(&self.account, None, 50) else {

            return;

        };

        self.media

            .register(

                page.entries

                    .iter()

                    .filter_map(|entry| self.data.events.get(&id(&entry.key)))

                    .flat_map(|event| event.media.clone())

                    .collect(),

            )

            .await;

        self.updates.send_replace(page);

    }

    /// writes the sync record, and the inbox with it when `dirty`, in one commit so a

    /// checkpoint never lands without the inbox changes it covers

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

        if self.dirty {

            self.advance()?;

        }

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

        if let Some(store) = &mut self.store {

            use super::state::json_len;

            use super::store::Write;

            let encode = |error| {

                Error::with_source(

                    ErrorKind::Store,

                    "encoding the mentions inbox failed",

                    error,

                )

            };

            let sync = serde_json::to_value(&self.sync).map_err(encode)?;

            let mut writes = Vec::new();

            if self.dirty {

                let available = store

                    .available(&["inbox", "sync"])

                    .saturating_sub(store.sealed_bound("sync", json_len(&sync)));

                let mut size = store.sealed_bound("inbox", json_len(&self.data));

                while size > available {

                    let Some(freed) = self.data.evict_one() else {

                        break;

                    };

                    size = size.saturating_sub(freed * 3);

                }

                writes.push(Write::Set(

                    "inbox".to_owned(),

                    serde_json::to_value(&self.data).map_err(encode)?,

                ));

            }

            writes.push(Write::Set("sync".to_owned(), sync));

            store.commit(store.revision(), writes).await?;

        }

        if self.dirty {

            self.dirty = false;

            self.publish().await;

        }

        Ok(())

    }

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

        self.data.revision = self

            .data

            .revision

            .checked_add(1)

            .ok_or_else(|| invalid("mentions revision exhausted"))?;

        self.data.bound(

            self.config.max_mentions_entries,

            self.config.max_mention_context_events,

            self.config.max_mention_candidates,

        );

        self.sync.inbox_revision = self.data.revision;

        Ok(())

    }

    async fn apply(&mut self, data: Option<Data>, sync: Option<SyncState>) -> Result<(), Error> {

        let dirty = self.dirty;

        let previous_data = data.map(|data| std::mem::replace(&mut self.data, data));

        let previous_sync = sync.map(|sync| std::mem::replace(&mut self.sync, sync));

        self.dirty |= previous_data.is_some();

        let result = self.flush().await;

        if result.is_err() {

            if let Some(data) = previous_data {

                self.data = data;

            }

            if let Some(sync) = previous_sync {

                self.sync = sync;

            }

            self.dirty = dirty;

        }

        result

    }

    async fn store_failed(&mut self, error: Error) {

        tracing::error!(%error, "saving the mentions inbox failed");

        self.data.gap(None, MentionGapReason::StoreFailed, None);

        self.dirty = true;

        self.flush_after = now().saturating_add(FLUSH_RETRY_MS);

        self.publish().await;

    }

    async fn set_repair_policy(&mut self, policy: HistoryRepairPolicy) -> Result<(), Error> {

        let mut next = self.data.clone();

        next.repair_policy = policy;

        self.apply(Some(next), None).await

    }

    async fn repair_resolved(&mut self, room_id: OwnedRoomId) -> Result<(), Error> {

        let mut sync = self.sync.clone();

        let coverage = sync.repair_coverage.entry(room_id.clone()).or_default();

        coverage.boundary_ms = None;

        coverage.last_confirmed_ms = Some(now());

        coverage.reliable = false;

        coverage.unresolved = false;

        sync.observed_rooms.insert(room_id);

        self.apply(None, Some(sync)).await

    }

    async fn command(&mut self, command: MentionsCommand) -> Result<CommandOutput, Error> {

        match command {

            MentionsCommand::Page { cursor, limit } => {

                let page = self.data.page(&self.account, cursor.as_deref(), limit)?;

                self.media

                    .register(

                        page.entries

                            .iter()

                            .filter_map(|entry| self.data.events.get(&id(&entry.key)))

                            .flat_map(|event| event.media.clone())

                            .collect(),

                    )

                    .await;

                Ok(CommandOutput::MentionsPage(page))

            }

            MentionsCommand::Context(key) => {

                let (context, media) = self.data.context(&key)?;

                self.media.register(media).await;

                Ok(CommandOutput::MentionContext(context))