macltok

use std::{

    collections::{BTreeMap, BTreeSet, HashMap},

    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::{

    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, SavedEvent, SyncState, Window, id, invalid, now};

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

use super::store::Opener;

use crate::{

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

    MentionGapReason, MentionKey, MentionReason, MentionsCommand, MentionsPage, MentionsUpdate,

    TimelineContent, TimelineItem,

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

    normalise::{MediaRegistration, MediaScope, normalise_observed, resolve_replies},

    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;

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

pub(crate) struct Unopened {

    pub(crate) error: Error,

    /// set when another attempt may succeed

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

    pub(crate) retry: Option<Opener>,

}

struct Loaded {

    store: Option<MentionsStore>,

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

    reopen: Option<Opener>,

    data: Data,

    sync: SyncState,

    rooms: BTreeSet<OwnedRoomId>,

    gaps: Vec<MentionGapReason>,

}

impl Loaded {

    fn empty(gaps: Vec<MentionGapReason>) -> Self {

        Self {

            store: None,

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

            reopen: None,

            data: Data::default(),

            sync: SyncState::default(),

            rooms: BTreeSet::new(),

            gaps,

        }

    }

}

#[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>>,

    ),

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

}

pub(crate) struct Service {

    client: Client,

    account: OwnedUserId,

    media: Arc<MediaService>,

    /// as many media as the inbox keeps events

    media_scope: MediaScope,

    store: Option<MentionsStore>,

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

    reopen: Option<Opener>,

    data: Data,

    sync: SyncState,

    rooms: BTreeSet<OwnedRoomId>,

    dirty: bool,

    rooms_dirty: bool,

    unpublished: bool,

    flush_after: u64,

    ignored: BTreeSet<OwnedUserId>,

    /// kept out of everything the inbox shows, while their entries stay saved for when the

    /// room stops being direct

    direct: BTreeSet<OwnedRoomId>,

    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 Work {

    fn rooms(&self) -> BTreeSet<OwnedRoomId> {

        let room = |key: &str| {

            key.split_once('\n')

                .and_then(|(room, _)| OwnedRoomId::try_from(room).ok())

        };

        match self {

            Self::Capture(key, ..) | Self::Candidate(key, _) => room(key).into_iter().collect(),

            Self::History(_, _, Ok(page)) => page

                .notifications

                .iter()

                .map(|notification| notification.room_id.clone())

                .collect(),

            Self::History(..) => BTreeSet::new(),

        }

    }

}

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 (reply, response) = oneshot::channel();

                let queued = handle.requests.send(Request::Batch(batch, reply));

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

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

                    Ok(Err(_)) => {

                        tracing::warn!("mentions service stopped, sync carries on");

                        return Ok(());

                    }

                    Err(_) => {

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

                        tracing::warn!("mentions queue stayed full, a response was dropped");

                        return Ok(());

                    }

                }

                // a queued batch still commits in order, so giving up here loses nothing

                // unless the process dies first, which the next start sees as a

                // discontinuity

                if matrix_sdk::timeout::timeout(response, HOOK_TIMEOUT)

                    .await

                    .is_err()

                {

                    tracing::warn!("mentions commit is slow, 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 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"))]

fn read_store(store: &MentionsStore) -> Result<(Data, SyncState, BTreeSet<OwnedRoomId>), Error> {

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

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

    let rooms = store.read("rooms")?.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, rooms))

}

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

async fn open_store(store: Option<Result<MentionsStore, Unopened>>) -> Loaded {

    let mut store = match store {

        None => return Loaded::empty(Vec::new()),

        Some(Err(unopened)) => {

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

            let mut loaded = Loaded::empty(vec![MentionGapReason::StoreFailed]);

            loaded.reopen = unopened.retry;

            return loaded;

        }

        Some(Ok(store)) => store,

    };

    let (data, sync, rooms) = match read_store(&store) {

        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 Loaded::empty(vec![MentionGapReason::StoreFailed]);

            }

            Default::default()

        }

    };

    let gaps = if store.was_reset() {

        vec![MentionGapReason::StoreReset]

    } else {

        Vec::new()

    };

    Loaded {

        store: Some(store),

        reopen: None,

        data,

        sync,

        rooms,

        gaps,

    }

}

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

async fn open_store(store: Option<Result<MentionsStore, Unopened>>) -> Loaded {

    match store {

        None => Loaded::empty(Vec::new()),

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

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

    }

}

impl Service {

    pub(crate) async fn new(

        client: Client,

        media: Arc<MediaService>,

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

        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 mut loaded = open_store(store).await;

        for reason in &loaded.gaps {

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

        }

        let direct = direct_rooms(&client);

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

        let media_scope = MediaScope::bounded(config.max_mention_context_events);

        media

            .register(

                &media_scope,

                page.entries

                    .iter()

                    .filter_map(|entry| loaded.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,

            media_scope,

            store: loaded.store,

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

            reopen: loaded.reopen,

            data: loaded.data,

            sync: loaded.sync,

            rooms: loaded.rooms,

            dirty: !loaded.gaps.is_empty(),

            rooms_dirty: false,

            unpublished: false,

            flush_after: 0,

            ignored: BTreeSet::new(),

            direct,

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

        self.unpublished = false;

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

            return;

        };

        self.media

            .register(

                &self.media_scope,

                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, with the inbox when `dirty` and observed rooms when

    /// `rooms_dirty`, in one commit so a checkpoint never lands without the 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::store::Write;

            let encode = |error| {

                Error::with_source(

                    ErrorKind::Store,

                    "encoding the mentions inbox failed",

                    error,

                )

            };

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

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

            let rooms = if self.rooms_dirty {

                Some(serde_json::value::to_raw_value(&self.rooms).map_err(encode)?)

            } else {

                None

            };

            let mut writes = Vec::new();

            if self.dirty {

                self.data

                    .coverage

                    .gaps

                    .retain(|gap| gap.reason != MentionGapReason::StoreFailed);

                let replaced: &[&str] = if rooms.is_some() {

                    &["inbox", "sync", "rooms"]

                } else {

                    &["inbox", "sync"]

                };

                let available = store

                    .available(replaced)

                    .saturating_sub(store.sealed_bound("sync", sync.get().len()))

                    .saturating_sub(

                        rooms

                            .as_ref()

                            .map_or(0, |rooms| store.sealed_bound("rooms", rooms.get().len())),

                    );

                let mut inbox = serde_json::value::to_raw_value(&self.data).map_err(encode)?;

                let mut size = store.sealed_bound("inbox", inbox.get().len());

                if size > available {

                    while size > available {

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

                            break;

                        };

                        size = size.saturating_sub(freed * 3);

                    }

                    inbox = serde_json::value::to_raw_value(&self.data).map_err(encode)?;

                }

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

            }

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

            if let Some(rooms) = rooms {

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

            }

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

        }

        self.rooms_dirty = false;

        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,

        );

        Ok(())

    }

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

        let dirty = std::mem::replace(&mut self.dirty, true);

        let previous = std::mem::replace(&mut self.data, data);

        let result = self.flush().await;

        if result.is_err() {

            self.data = previous;

            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;

    }