macltok

use std::{

    collections::{BTreeMap, BTreeSet},

    sync::Arc,

    time::Duration,

};

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

use matrix_sdk::{

    Client, RoomState,

    sliding_sync::{ResponseCommitBatch, ResponseCommitHook},

};

use ruma::{

    OwnedRoomId, OwnedUserId,

    events::{AnySyncTimelineEvent, room::encrypted::OriginalSyncRoomEncryptedEvent},

    serde::Raw,

};

use serde_json::Value;

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

use super::state::{Candidate, Data, SavedEvent, Window, id, invalid, now};

use crate::{

    CommandOutput, EngineConfig, Error, ErrorKind, Event, MentionCaptureState, MentionEntry,

    MentionGap, 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;

#[derive(Clone)]

pub(crate) struct Handle(mpsc::Sender<Request>);

enum Request {

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

    Command(

        MentionsCommand,

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

    ),

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

}

pub(crate) struct Service {

    client: Client,

    account: OwnedUserId,

    media: Arc<MediaService>,

    store: Option<MentionsStore>,

    data: Data,

    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 {

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

        let handle = self.clone();

        ResponseCommitHook::new(move |batch| {

            let handle = handle.clone();

            async move {

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

                handle

                    .0

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

                    .await

                    .map_err(|_| "mentions service stopped".to_owned())?;

                response

                    .await

                    .map_err(|_| "mentions commit acknowledgement dropped".to_owned())?

                    .map_err(|error| error.to_string())

            }

        })

    }

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

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

        self.0

            .send(Request::Command(command, 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.0.send(Request::Stop(reply)).await.is_ok() {

            let _ = response.await;

        }

    }

}

impl Service {

    pub(crate) async fn new(

        client: Client,

        media: Arc<MediaService>,

        store: Option<MentionsStore>,

        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();

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

        let data: Data = store

            .as_ref()

            .map(|store| store.read("inbox"))

            .transpose()?

            .flatten()

            .unwrap_or_default();

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

        let data = Data::default();

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

        if store

            .as_ref()

            .is_some_and(|store| store.revision() != data.revision)

        {

            return Err(Error::new(

                ErrorKind::Store,

                "mentions revision does not match its checkpoint",

            ));

        }

        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));

        Ok((

            Handle(sender),

            Self {

                client,

                account,

                media,

                store,

                data,

                config,

                requests,

                updates,

                recovery_budget: 20,

                latest_scan: true,

                recovery_retry: 0,

            },

            receiver,

        ))

    }

    async fn save(&mut self, mut next: Data) -> Result<(), Error> {

        next.revision = self

            .data

            .revision

            .checked_add(1)

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

        next.bound(

            self.config.max_mentions_entries,

            self.config.max_mention_context_events,

            self.config.max_mention_candidates,

        );

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

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

            loop {

                let value = serde_json::to_value(&next).map_err(|error| {

                    Error::with_source(

                        ErrorKind::Store,

                        "encoding the mentions inbox failed",

                        error,

                    )

                })?;

                match store

                    .commit(

                        self.data.revision,

                        vec![super::store::Write::Set("inbox".to_owned(), value)],

                    )

                    .await

                {

                    Ok(_) => break,

                    Err(error)

                        if error.message == "mentions storage limit exceeded"

                            && next.evict_one() => {}

                    Err(error) => return Err(error),

                }

            }

        }

        self.data = next;

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

        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);

        Ok(())

    }

    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))

            }

            MentionsCommand::RetryContext(key) => {

                let room = self

                    .client

                    .get_room(&key.room_id)

                    .filter(|room| room.state() == RoomState::Joined)

                    .ok_or_else(|| {

                        Error::new(

                            ErrorKind::Forbidden,

                            "saved context cannot be refetched after leaving the room",

                        )

                    })?;

                drop(room);

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

                let entry = next

                    .entries

                    .get_mut(&id(&key))

                    .ok_or_else(|| Error::new(ErrorKind::EventNotFound, "mention was not found"))?;

                entry.capture = MentionCaptureState::Pending;

                entry.revision = self.data.revision + 1;

                next.windows.entry(id(&key)).or_insert(Window {

                    before: Vec::new(),

                    target: Some(id(&key)),

                    after: Vec::new(),

                    start: None,

                    end: None,

                });

                self.save(next).await?;

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

                self.media.register(media).await;

                Ok(CommandOutput::MentionContext(context))

            }

            MentionsCommand::SetRead { keys, read } => {

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

                for key in keys {

                    if let Some(entry) = next.entries.get_mut(&id(&key)) {

                        entry.read = read;

                        entry.revision = self.data.revision + 1;

                    }

                }

                self.save(next).await?;

                Ok(CommandOutput::Done)

            }

            MentionsCommand::ReadThrough(boundary) => {

                if boundary.account != self.account || boundary.revision > self.data.revision {

                    return Err(invalid("invalid mentions read boundary"));

                }

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

                for entry in next

                    .entries

                    .values_mut()

                    .filter(|entry| entry.discovery_revision <= boundary.revision)

                {

                    entry.read = true;

                    entry.revision = self.data.revision + 1;

                }

                self.save(next).await?;

                Ok(CommandOutput::Done)

            }

            MentionsCommand::Recover => {

                self.recovery_budget = 20;

                self.latest_scan = true;

                self.recovery_retry = 0;

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

                next.coverage.recovering = true;

                next.coverage.historical_supported = None;

                self.save(next).await?;

                Ok(CommandOutput::Done)

            }

        }

    }

    fn gap(data: &mut Data, room_id: Option<OwnedRoomId>, reason: &str, token: Option<String>) {

        if !data

            .coverage

            .gaps

            .iter()

            .any(|gap| gap.room_id == room_id && gap.reason == reason && gap.token == token)

        {

            data.coverage.gaps.push(MentionGap {

                room_id,

                reason: reason.to_owned(),

                token,

            });

            if data.coverage.gaps.len() > 128 {

                data.coverage.gaps.remove(0);

                data.coverage.truncated = true;

            }

        }

    }

    async fn batch(&mut self, batch: ResponseCommitBatch) -> Result<(), Error> {

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

        if next.checkpoints.get(&batch.connection_id) == Some(&batch.next_pos) {

            return Ok(());

        }

        if next.checkpoints.get(&batch.connection_id) != batch.previous_pos.as_ref() {

            Self::gap(

                &mut next,

                None,

                "sync checkpoint continuity is unknown",

                None,

            );

        }

        for (room_id, raw) in &batch.raw_rooms {

            if raw.limited || !next.observed_rooms.contains(room_id) {

                Self::gap(

                    &mut next,

                    Some(room_id.clone()),

                    if raw.limited {

                        "limited sync timeline"

                    } else {

                        "initial room history is incomplete"

                    },

                    raw.prev_batch.clone(),

                );

            }

        }

        let room_ids: BTreeSet<_> = batch

            .raw_rooms

            .keys()

            .chain(batch.response.rooms.joined.keys())

            .cloned()

            .collect();

        for room_id in &room_ids {

            if !self

                .client

                .get_room(room_id)

                .is_some_and(|room| room.state() == RoomState::Joined)

            {

                continue;

            }

            next.observed_rooms.insert(room_id.clone());

            let mut observed = BTreeMap::new();

            if let Some(raw) = batch.raw_rooms.get(room_id) {

                for event in &raw.timeline {

                    let value: Value =

                        serde_json::from_str(event.json().get()).map_err(|error| {

                            Error::with_source(ErrorKind::Invalid, "invalid observed event", error)

                        })?;

                    if let Some(event_id) = value["event_id"].as_str() {

                        observed.insert(event_id.to_owned(), value);

                    }

                }

            }

            if let Some(update) = batch.response.rooms.joined.get(room_id) {

                for event in &update.timeline.events {

                    let value: Value =

                        serde_json::from_str(event.raw().json().get()).map_err(|error| {

                            Error::with_source(ErrorKind::Invalid, "invalid observed event", error)

                        })?;

                    if let Some(event_id) = value["event_id"].as_str() {

                        observed.insert(event_id.to_owned(), value);

                    }

                }

            }

            let mut observed: Vec<_> = observed.into_values().collect();

            observed.sort_by_key(|value| value["origin_server_ts"].as_u64().unwrap_or(0));

            for value in observed {

                self.observe(&mut next, room_id, value).await?;

            }

        }

        for room_id in batch.response.rooms.left.keys() {

            Self::gap(

                &mut next,

                Some(room_id.clone()),

                "room left, retained context is read only",

                None,

            );

            next.candidates

                .retain(|_, candidate| &candidate.room_id != room_id);

            for entry in next

                .entries

                .values_mut()

                .filter(|entry| &entry.key.room_id == room_id)

            {

                entry.capture = MentionCaptureState::Unavailable;

                entry.revision = next.revision + 1;

            }

        }

        let reevaluate: Vec<_> = next

            .entries

            .values()

            .filter(|entry| room_ids.contains(&entry.key.room_id))

            .filter_map(|entry| {

                next.events

                    .get(&id(&entry.key))

                    .and_then(|event| event.raw.clone())

                    .map(|raw| (entry.key.room_id.clone(), raw))

            })

            .collect();

        for (room_id, raw) in reevaluate {

            self.observe(&mut next, &room_id, raw).await?;

        }

        self.filter_ignored(&mut next).await;

        next.checkpoints.insert(batch.connection_id, batch.next_pos);

        self.save(next).await

    }

    async fn filter_ignored(&self, next: &mut Data) {

        let mut removed = Vec::new();

        for (key, entry) in &next.entries {

            if let TimelineItem::Event(event) = &entry.event {

                if self.client.is_user_ignored(&event.sender.user_id).await {

                    removed.push(key.clone());

                }

            }

        }

        for key in removed {

            next.entries.remove(&key);

            next.windows.remove(&key);

        }

        next.collect_events();

    }

    async fn observe(

        &self,

        data: &mut Data,

        room_id: &OwnedRoomId,

        mut value: Value,

    ) -> Result<(), Error> {

        let Some(event_id) = value["event_id"]

            .as_str()

            .and_then(|value| value.parse().ok())

        else {

            return Ok(());