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

    OwnedEventId, OwnedRoomId, OwnedUserId,

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

    serde::Raw,

};

use serde_json::Value;

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

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

use crate::{

    CommandOutput, EngineConfig, Error, ErrorKind, Event, HistoryRepairPolicy, 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>>,

    ),

    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,

    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 repair_policy(&self) -> Result<HistoryRepairPolicy, Error> {

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

        self.0

            .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.0

            .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.0

            .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.0

            .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.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 set_repair_policy(&mut self, policy: HistoryRepairPolicy) -> Result<(), Error> {

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

        next.repair_policy = policy;

        self.save(next).await

    }

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

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

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

        coverage.boundary_ms = None;

        coverage.last_confirmed_ms = Some(now());

        coverage.reliable = false;

        coverage.unresolved = false;

        next.observed_rooms.insert(room_id);

        self.save(next).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))

            }

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

        }

        let observed_at_ms = now();

        let checkpoint_discontinuity =

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

        if checkpoint_discontinuity {

            Self::gap(

                &mut next,

                None,

                "sync checkpoint continuity is unknown",

                None,

            );

            let observed_rooms: Vec<_> = next.observed_rooms.iter().cloned().collect();

            for room_id in observed_rooms {

                Self::mark_unresolved(&mut next, &room_id, observed_at_ms);

            }

        }

        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;

            }

            let previously_observed = next.observed_rooms.contains(room_id);

            let limited = batch.raw_rooms.get(room_id).is_some_and(|raw| raw.limited);

            if checkpoint_discontinuity || limited || !previously_observed {

                Self::mark_unresolved(&mut next, room_id, observed_at_ms);

            } else {

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

                if !coverage.unresolved {

                    coverage.last_confirmed_ms = Some(observed_at_ms);

                }

                coverage.reliable = false;

            }

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

            let mut observed = BTreeMap::new();

            let mut latest_anchor = None;

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

                        latest_anchor = event_id.parse::<OwnedEventId>().ok();

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