macltok

use std::{collections::HashSet, future::Future, panic::AssertUnwindSafe, sync::Arc};

use futures_util::{FutureExt, StreamExt};

use matrix_sdk::SendOutsideWasm;

use matrix_sdk::{

    room::edit::EditedContent,

    room::reply::{EnforceThread, Reply as SdkReply},

    ruma::{

        api::client::receipt::create_receipt::v3::ReceiptType,

        events::room::message::{

            AddMentions, Relation, ReplyWithinThread, RoomMessageEventContentWithoutRelation,

        },

    },

};

use matrix_sdk_ui::timeline::{Timeline, TimelineEventItemId};

use ruma::{OwnedEventId, OwnedUserId};

use tokio::sync::mpsc;

use super::executor::{EngineBoxFuture, JoinHandle, TaskSet, spawn};

use crate::{

    Error, ErrorKind, Event, EventRef, HistoryState, MessageContent, MessageEditContent,

    MessageMentions, RichText, TimelineCommand, TimelineId,

    engine::MediaService,

    normalise::{MediaScope, normalise_timeline, normalise_timeline_diffs},

    transport::{EventSender, Responder, SendEventError},

    upload::sdk_mentions,

};

const PAGINATION_QUEUE: usize = 32;

/// long-lived tasks inside a room actor report panics here so the actor can die loudly

#[derive(Clone)]

pub(super) struct Supervisor(mpsc::UnboundedSender<&'static str>);

impl Supervisor {

    pub(super) fn new() -> (Self, mpsc::UnboundedReceiver<&'static str>) {

        let (sender, receiver) = mpsc::unbounded_channel();

        (Self(sender), receiver)

    }

    pub(super) fn spawn(

        &self,

        label: &'static str,

        task: impl Future<Output = ()> + SendOutsideWasm + 'static,

    ) -> JoinHandle<()> {

        let panics = self.0.clone();

        spawn(async move {

            if AssertUnwindSafe(task).catch_unwind().await.is_err() {

                let _ = panics.send(label);

            }

        })

    }

}

/// aborts its task when dropped, so an actor that dies mid-flight takes its helpers with it

pub(super) struct TaskGuard(pub(super) JoinHandle<()>);

impl Drop for TaskGuard {

    fn drop(&mut self) {

        self.0.abort();

    }

}

/// what a timeline session needs from the room that owns it

#[derive(Clone)]

pub(super) struct SessionContext {

    pub(super) events: EventSender,

    pub(super) media: Arc<MediaService>,

    pub(super) mentions: crate::mentions::service::Handle,

    pub(super) current_user_id: Option<OwnedUserId>,

    pub(super) initial_history_limit: u16,

    pub(super) supervisor: Supervisor,

}

pub(super) struct TimelineSession {

    pub(super) id: TimelineId,

    timeline: Arc<Timeline>,

    context: SessionContext,

    scope: MediaScope,

    watcher: JoinHandle<()>,

    pagination: mpsc::Sender<PaginationRequest>,

    pagination_worker: JoinHandle<()>,

    member_fetch: JoinHandle<()>,

}

struct PaginationRequest {

    limit: u16,

    reply: Option<Responder>,

}

impl TimelineSession {

    pub(super) async fn start(

        id: TimelineId,

        timeline: Timeline,

        context: SessionContext,

    ) -> Result<Self, SendEventError> {

        let timeline = Arc::new(timeline);

        let scope = MediaScope::new();

        let watcher = start_watcher(&id, &timeline, &context, &scope).await?;

        let (pagination, pagination_worker) = start_pagination_worker(&id, &timeline, &context);

        // under lazy-loaded members the sdk leaves sender profiles unavailable until someone

        // asks, and a /members fetched any other way never reaches the timeline's items

        let member_fetch = context.supervisor.spawn("member fetch", {

            let timeline = timeline.clone();

            async move { timeline.fetch_members().await }

        });

        let session = Self {

            id,

            timeline,

            scope,

            watcher,

            pagination,

            pagination_worker,

            member_fetch,

            context,

        };

        if session.context.initial_history_limit > 0 && !matches!(session.id, TimelineId::Room(_)) {

            session.paginate(session.context.initial_history_limit, None);

        }

        Ok(session)

    }

    /// swaps in a fresh watcher. the old one is dead before the reset goes out,

    /// so no stale diff can land after it

    pub(super) async fn resync(&mut self) -> Result<(), Error> {

        self.watcher.abort();

        let _ = (&mut self.watcher).await;

        self.watcher = start_watcher(&self.id, &self.timeline, &self.context, &self.scope)

            .await

            .map_err(|_| Error::new(ErrorKind::Timeline, "timeline reset could not be sent"))?;

        Ok(())

    }

    pub(super) async fn publish_history(&mut self) -> Result<(), Error> {

        self.timeline

            .flush_history()

            .await

            .map_err(timeline_error)?;

        self.resync().await?;

        self.context

            .events

            .publish_barrier()

            .await

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

    }

    pub(super) fn paginate(&self, limit: u16, reply: Option<Responder>) {

        if let Err(error) = self.pagination.try_send(PaginationRequest { limit, reply }) {

            let (request, error) = match error {

                mpsc::error::TrySendError::Full(request) => (

                    request,

                    Error::new(ErrorKind::QueueFull, "pagination queue is full"),

                ),

                mpsc::error::TrySendError::Closed(request) => (

                    request,

                    Error::new(ErrorKind::TimelineNotOpen, "timeline is closing"),

                ),

            };

            if let Some(reply) = request.reply {

                reply.done(Err(error));

            }

        }

    }

    async fn stop_without_close(&mut self) {

        self.member_fetch.abort();

        self.pagination_worker.abort();

        self.watcher.abort();

        let _ = (&mut self.member_fetch).await;

        let _ = (&mut self.pagination_worker).await;

        let _ = (&mut self.watcher).await;

        // after the watcher, or a last diff could register media for a timeline that's gone

        self.context.media.release(&self.scope).await;

    }

    /// pending paginations answer `Interrupted` as their responders drop

    pub(super) async fn close(mut self) {

        self.stop_without_close().await;

        let _ = self

            .context

            .events

            .send(Event::TimelineClosed {

                timeline_id: self.id.clone(),

            })

            .await;

    }

    pub(super) async fn execute(

        &self,

        command: TimelineCommand,

        reply: Responder,

        work: &mut TaskSet<()>,

    ) {

        let timeline = self.timeline.clone();

        let thread_root = self.id.thread_root().cloned();

        let background: EngineBoxFuture<'static, Result<(), Error>> = match command {

            TimelineCommand::Paginate { limit, .. } => return self.paginate(limit, Some(reply)),

            // network-bound, so it runs beside the actor instead of queueing the room behind it

            TimelineCommand::Redact { target, reason, .. } => Box::pin(async move {

                timeline

                    .redact(&item_id(target), reason.as_deref())

                    .await

                    .map_err(timeline_error)

            }),

            TimelineCommand::ToggleReaction {

                target,

                key,

                shortcode,

                ..

            } => Box::pin(async move {

                timeline

                    .toggle_reaction_with_shortcode(&item_id(target), &key, shortcode.as_deref())

                    .await

                    .map(drop)

                    .map_err(timeline_error)

            }),

            TimelineCommand::MarkRead { event_id, .. } => Box::pin(async move {

                let receipt_failed =

                    |source| Error::from_sdk("sending read receipt failed", source);

                timeline

                    .send_single_receipt(ReceiptType::Read, event_id.clone())

                    .await

                    .map_err(receipt_failed)?;

                if thread_root.is_none() {

                    timeline

                        .send_single_receipt(ReceiptType::FullyRead, event_id)

                        .await

                        .map_err(receipt_failed)?;

                }

                Ok(())

            }),

            command => return reply.done(self.execute_serial(command).await),

        };

        work.spawn(async move { reply.done(background.await) });

    }

    /// sends stay on the actor so they hit the send queue in the order they were asked for

    async fn execute_serial(&self, command: TimelineCommand) -> Result<(), Error> {

        let timeline = &self.timeline;

        match command {

            TimelineCommand::Send { content, .. } => {

                let body = match &content {

                    MessageContent::Text { content, .. }

                    | MessageContent::Emote { content, .. }

                    | MessageContent::Notice { content, .. }

                    | MessageContent::Reply { content, .. } => &content.body,

                };

                if body.trim().is_empty() {

                    return Err(Error::new(ErrorKind::EmptyBody, "message body is empty"));

                }

                let content = match content {

                    MessageContent::Text { content, mentions } => {

                        message_content(content, mentions, Msgtype::Text).with_relation(None)

                    }

                    MessageContent::Emote { content, mentions } => {

                        message_content(content, mentions, Msgtype::Emote).with_relation(None)

                    }

                    MessageContent::Notice { content, mentions } => {

                        message_content(content, mentions, Msgtype::Notice).with_relation(None)

                    }

                    MessageContent::Reply {

                        event_id,

                        content,

                        mentions,

                        notify_replied_to,

                    } => {

                        self.reply_content(event_id, content, mentions, notify_replied_to)

                            .await?

                    }

                };

                timeline

                    .send(content.into())

                    .await

                    .map(drop)

                    .map_err(timeline_error)

            }

            TimelineCommand::Edit {

                target, content, ..

            } => {

                let (content, mentions, msgtype) = match content {

                    MessageEditContent::Text { content, mentions } => {

                        (content, mentions, Msgtype::Text)

                    }

                    MessageEditContent::Emote { content, mentions } => {

                        (content, mentions, Msgtype::Emote)

                    }

                };

                if content.body.trim().is_empty() {

                    return Err(Error::new(ErrorKind::EmptyBody, "message body is empty"));

                }

                timeline

                    .edit(

                        &item_id(target),

                        EditedContent::RoomMessage(message_content(content, mentions, msgtype)),

                    )

                    .await

                    .map_err(timeline_error)

            }

            _ => Err(Error::new(

                ErrorKind::Other,

                "timeline command was misrouted",

            )),

        }

    }

    async fn reply_content(

        &self,

        event_id: OwnedEventId,

        content: RichText,

        mentions: MessageMentions,

        notify_replied_to: bool,

    ) -> Result<matrix_sdk::ruma::events::room::message::RoomMessageEventContent, Error> {

        let thread_root = self.id.thread_root();

        let content = self

            .timeline

            .room()

            .make_reply_event(

                message_content(content, mentions, Msgtype::Text),

                SdkReply {

                    event_id,

                    enforce_thread: if thread_root.is_some() {

                        EnforceThread::Threaded(ReplyWithinThread::Yes)

                    } else {

                        EnforceThread::MaybeThreaded

                    },

                    add_mentions: if notify_replied_to {

                        AddMentions::Yes

                    } else {

                        AddMentions::No

                    },

                },

            )

            .await

            .map_err(|source| {

                Error::with_source(ErrorKind::Timeline, "reply operation failed", source)

            })?;

        if let Some(thread_root) = thread_root

            && !matches!(

                content.relates_to.as_ref(),

                Some(Relation::Thread(thread)) if &thread.event_id == thread_root

            )

        {

            return Err(Error::new(

                ErrorKind::Invalid,

                "reply target does not belong to the selected thread",

            ));

        }

        Ok(content)

    }

}

fn start_pagination_worker(

    id: &TimelineId,

    timeline: &Arc<Timeline>,

    context: &SessionContext,

) -> (mpsc::Sender<PaginationRequest>, JoinHandle<()>) {

    let (pagination, requests) = mpsc::channel(PAGINATION_QUEUE);

    let worker = context.supervisor.spawn(

        "pagination worker",

        run_pagination(

            requests,

            timeline.clone(),

            context.events.clone(),

            id.clone(),

        ),

    );

    (pagination, worker)

}

impl Drop for TimelineSession {

    fn drop(&mut self) {

        self.watcher.abort();

        self.pagination_worker.abort();

        self.member_fetch.abort();

    }

}

fn item_id(target: EventRef) -> TimelineEventItemId {

    match target {

        EventRef::Remote(event_id) => TimelineEventItemId::EventId(event_id),

        EventRef::Local(transaction_id) => TimelineEventItemId::TransactionId(transaction_id),

    }

}

/// the sdk leaves a reply's target unloaded until someone asks, so the watcher asks, once per event

struct ReplyFetcher {

    timeline: Arc<Timeline>,

    asked: HashSet<OwnedEventId>,

    fetches: TaskSet<()>,

}

impl ReplyFetcher {

    fn fetch(&mut self, missing: Vec<OwnedEventId>) {

        for event_id in missing {

            if self.asked.insert(event_id.clone()) {

                let timeline = self.timeline.clone();

                self.fetches.spawn(async move {

                    if let Err(error) = timeline.fetch_details_for_event(&event_id).await {

                        tracing::debug!(%event_id, %error, "fetching reply details failed");

                    }

                });

            }

        }

    }

}

async fn start_watcher(

    id: &TimelineId,

    timeline: &Arc<Timeline>,

    context: &SessionContext,

    scope: &MediaScope,

) -> Result<JoinHandle<()>, SendEventError> {

    let (items, mut stream) = timeline.subscribe().await;

    let current_user_id = context.current_user_id.clone();

    let normalised = normalise_timeline(items.iter(), current_user_id.as_deref(), scope);

    context.media.register(normalised.media).await;

    context

        .events

        .send(Event::TimelineReset {

            timeline_id: id.clone(),

            items: normalised.items,

        })

        .await?;

    let mut replies = ReplyFetcher {

        timeline: timeline.clone(),

        asked: HashSet::new(),

        fetches: TaskSet::new(),

    };

    replies.fetch(normalised.missing_replies);

    let id = id.clone();

    let events = context.events.clone();

    let media = context.media.clone();

    let scope = scope.clone();

    Ok(context.supervisor.spawn("timeline watcher", async move {

        loop {

            let diffs = tokio::select! {

                diffs = stream.next() => match diffs {

                    Some(diffs) => diffs,

                    None => return,

                },

                Some(_) = replies.fetches.join_next(), if !replies.fetches.is_empty() => continue,

            };

            let normalised = normalise_timeline_diffs(diffs, current_user_id.as_deref(), &scope);

            media.register(normalised.media).await;

            replies.fetch(normalised.missing_replies);

            let event = Event::TimelineChanged {

                timeline_id: id.clone(),

                diffs: normalised.diffs,

            };

            if events.send(event).await.is_err() {

                return;

            }

        }

    }))

}

async fn run_pagination(

    mut requests: mpsc::Receiver<PaginationRequest>,

    timeline: Arc<Timeline>,

    events: EventSender,

    timeline_id: TimelineId,

) {

    while let Some(PaginationRequest { limit, reply }) = requests.recv().await {

        let history = |state| Event::HistoryChanged {

            timeline_id: timeline_id.clone(),

            state,

        };

        if events.send(history(HistoryState::Loading)).await.is_err() {

            return;

        }

        let result = timeline

            .paginate_backwards(limit.max(1))

            .await

            .map_err(timeline_error);

        let state = match &result {

            Ok(true) => HistoryState::StartReached,

            Ok(false) => HistoryState::Idle,

            Err(error) => HistoryState::Failed(error.to_string()),

        };

        let _ = events.send(history(state)).await;

        if let Some(reply) = reply {